From 23353bbd1bbe49a2a95617f69933ee99bf7d04f5 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Fri, 4 Aug 2023 08:52:05 -0700 Subject: [PATCH] Improvements to `pythonTerminalEnvVarActivation` experiment (#21751) --- src/client/interpreter/activation/service.ts | 51 ++++++++-- .../terminalEnvVarCollectionService.ts | 58 ++++++----- src/client/interpreter/activation/types.ts | 2 + ...rminalEnvVarCollectionService.unit.test.ts | 97 ++++++++----------- 4 files changed, 118 insertions(+), 90 deletions(-) diff --git a/src/client/interpreter/activation/service.ts b/src/client/interpreter/activation/service.ts index e5da57227b19..4364cc825f78 100644 --- a/src/client/interpreter/activation/service.ts +++ b/src/client/interpreter/activation/service.ts @@ -19,8 +19,8 @@ import { ICurrentProcess, IDisposable, Resource } from '../../common/types'; import { sleep } from '../../common/utils/async'; import { InMemoryCache } from '../../common/utils/cacheUtils'; import { OSType } from '../../common/utils/platform'; -import { IEnvironmentVariablesProvider } from '../../common/variables/types'; -import { EnvironmentType, PythonEnvironment } from '../../pythonEnvironments/info'; +import { EnvironmentVariables, IEnvironmentVariablesProvider } from '../../common/variables/types'; +import { EnvironmentType, PythonEnvironment, virtualEnvTypes } from '../../pythonEnvironments/info'; import { sendTelemetryEvent } from '../../telemetry'; import { EventName } from '../../telemetry/constants'; import { IInterpreterService } from '../contracts'; @@ -38,6 +38,7 @@ import { Conda } from '../../pythonEnvironments/common/environmentManagers/conda import { StopWatch } from '../../common/utils/stopWatch'; import { identifyShellFromShellPath } from '../../common/terminal/shellDetectors/baseShellDetector'; import { getSearchPathEnvVarNames } from '../../common/utils/exec'; +import { cache } from '../../common/utils/decorators'; const ENVIRONMENT_PREFIX = 'e8b39361-0157-4923-80e1-22d70d46dee6'; const CACHE_DURATION = 10 * 60 * 1000; @@ -154,11 +155,11 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi } // Cache only if successful, else keep trying & failing if necessary. - const cache = new InMemoryCache(CACHE_DURATION); + const memCache = new InMemoryCache(CACHE_DURATION); return this.getActivatedEnvironmentVariablesImpl(resource, interpreter, allowExceptions, shell) .then((vars) => { - cache.data = vars; - this.activatedEnvVariablesCache.set(cacheKey, cache); + memCache.data = vars; + this.activatedEnvVariablesCache.set(cacheKey, memCache); sendTelemetryEvent( EventName.PYTHON_INTERPRETER_ACTIVATION_ENVIRONMENT_VARIABLES, stopWatch.elapsedTime, @@ -176,6 +177,35 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi }); } + @cache(-1, true) + public async getProcessEnvironmentVariables(resource: Resource, shell?: string): Promise { + // Try to get the process environment variables using Python by printing variables, that can be little different + // from `process.env` and is preferred when calculating diff. + const globalInterpreters = this.interpreterService + .getInterpreters() + .filter((i) => !virtualEnvTypes.includes(i.envType)); + const interpreterPath = + globalInterpreters.length > 0 && globalInterpreters[0] ? globalInterpreters[0].path : 'python'; + try { + const [args, parse] = internalScripts.printEnvVariables(); + args.forEach((arg, i) => { + args[i] = arg.toCommandArgumentForPythonExt(); + }); + const command = `${interpreterPath} ${args.join(' ')}`; + const processService = await this.processServiceFactory.create(resource); + const result = await processService.shellExec(command, { + shell, + timeout: ENVIRONMENT_TIMEOUT, + maxBuffer: 1000 * 1000, + throwOnStdErr: false, + }); + const returnedEnv = this.parseEnvironmentOutput(result.stdout, parse); + return returnedEnv ?? process.env; + } catch (ex) { + return process.env; + } + } + public async getEnvironmentActivationShellCommands( resource: Resource, interpreter?: PythonEnvironment, @@ -231,7 +261,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi ); traceVerbose(`Activation Commands received ${activationCommands} for shell ${shellInfo.shell}`); if (!activationCommands || !Array.isArray(activationCommands) || activationCommands.length === 0) { - if (interpreter?.envType === EnvironmentType.Venv) { + if (interpreter && [EnvironmentType.Venv, EnvironmentType.Pyenv].includes(interpreter?.envType)) { const key = getSearchPathEnvVarNames()[0]; if (env[key]) { env[key] = `${path.dirname(interpreter.path)}${path.delimiter}${env[key]}`; @@ -247,7 +277,14 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi const activationCommand = fixActivationCommands(activationCommands).join(' && '); // In order to make sure we know where the environment output is, // put in a dummy echo we can look for - command = `${activationCommand} && echo '${ENVIRONMENT_PREFIX}' && python ${args.join(' ')}`; + const commandSeparator = [TerminalShellType.powershell, TerminalShellType.powershellCore].includes( + shellInfo.shellType, + ) + ? ';' + : '&&'; + command = `${activationCommand} ${commandSeparator} echo '${ENVIRONMENT_PREFIX}' ${commandSeparator} python ${args.join( + ' ', + )}`; } // Make sure python warnings don't interfere with getting the environment. However diff --git a/src/client/interpreter/activation/terminalEnvVarCollectionService.ts b/src/client/interpreter/activation/terminalEnvVarCollectionService.ts index 85b393425836..f4c29f03707d 100644 --- a/src/client/interpreter/activation/terminalEnvVarCollectionService.ts +++ b/src/client/interpreter/activation/terminalEnvVarCollectionService.ts @@ -33,6 +33,7 @@ import { defaultShells } from './service'; import { IEnvironmentActivationService } from './types'; import { EnvironmentType } from '../../pythonEnvironments/info'; import { getSearchPathEnvVarNames } from '../../common/utils/exec'; +import { EnvironmentVariables } from '../../common/variables/types'; @injectable() export class TerminalEnvVarCollectionService implements IExtensionActivationService { @@ -45,7 +46,10 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ private registeredOnce = false; - private previousEnvVars = _normCaseKeys(process.env); + /** + * Carries default environment variables for the currently selected shell. + */ + private processEnvVars: EnvironmentVariables | undefined; constructor( @inject(IPlatformService) private readonly platform: IPlatformService, @@ -90,6 +94,7 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ this.applicationEnvironment.onDidChangeShell( async (shell: string) => { this.showProgress(); + this.processEnvVars = undefined; // Pass in the shell where known instead of relying on the application environment, because of bug // on VSCode: https://github.com/microsoft/vscode/issues/160694 await this._applyCollection(undefined, shell).ignoreErrors(); @@ -106,6 +111,9 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ public async _applyCollection(resource: Resource, shell = this.applicationEnvironment.shell): Promise { const workspaceFolder = this.getWorkspaceFolder(resource); const settings = this.configurationService.getSettings(resource); + const envVarCollection = this.getEnvironmentVariableCollection(workspaceFolder); + // Clear any previously set env vars from collection. + envVarCollection.clear(); if (!settings.terminal.activateEnvironment) { traceVerbose('Activating environments in terminal is disabled for', resource?.fsPath); return; @@ -116,7 +124,6 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ undefined, shell, ); - const envVarCollection = this.getEnvironmentVariableCollection(workspaceFolder); if (!env) { const shellType = identifyShellFromShellPath(shell); const defaultShell = defaultShells[this.platform.osType]; @@ -126,32 +133,38 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ await this._applyCollection(resource, defaultShell?.shell); return; } - envVarCollection.clear(); - this.previousEnvVars = _normCaseKeys(process.env); + this.processEnvVars = undefined; return; } - const previousEnv = this.previousEnvVars; - this.previousEnvVars = env; + if (!this.processEnvVars) { + this.processEnvVars = await this.environmentActivationService.getProcessEnvironmentVariables( + resource, + shell, + ); + } + const processEnv = this.processEnvVars; Object.keys(env).forEach((key) => { + if (shouldSkip(key)) { + return; + } const value = env[key]; - const prevValue = previousEnv[key]; + const prevValue = processEnv[key]; if (prevValue !== value) { if (value !== undefined) { + if (key === 'PS1') { + traceVerbose(`Prepending environment variable ${key} in collection with ${value}`); + envVarCollection.prepend(key, value, { + applyAtShellIntegration: true, + applyAtProcessCreation: false, + }); + return; + } traceVerbose(`Setting environment variable ${key} in collection to ${value}`); envVarCollection.replace(key, value, { applyAtShellIntegration: true }); - } else { - traceVerbose(`Clearing environment variable ${key} from collection`); - envVarCollection.delete(key); } } }); - Object.keys(previousEnv).forEach((key) => { - // If the previous env var is not in the current env, clear it from collection. - if (!(key in env)) { - traceVerbose(`Clearing environment variable ${key} from collection`); - envVarCollection.delete(key); - } - }); + const displayPath = this.pathUtils.getDisplayName(settings.pythonPath, workspaceFolder?.uri.fsPath); const description = new MarkdownString(`${Interpreters.activateTerminalDescription} \`${displayPath}\``); envVarCollection.description = description; @@ -224,13 +237,6 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ } } -export function _normCaseKeys(env: NodeJS.ProcessEnv): NodeJS.ProcessEnv { - const result: NodeJS.ProcessEnv = {}; - Object.keys(env).forEach((key) => { - // `os.environ` script used to get env vars normalizes keys to upper case: - // https://github.com/python/cpython/issues/101754 - // So convert `process.env` keys to upper case to match. - result[key.toUpperCase()] = env[key]; - }); - return result; +function shouldSkip(env: string) { + return ['_', 'SHLVL'].includes(env); } diff --git a/src/client/interpreter/activation/types.ts b/src/client/interpreter/activation/types.ts index d8e4ae16dbca..e00ef9b62b3f 100644 --- a/src/client/interpreter/activation/types.ts +++ b/src/client/interpreter/activation/types.ts @@ -4,10 +4,12 @@ 'use strict'; import { Resource } from '../../common/types'; +import { EnvironmentVariables } from '../../common/variables/types'; import { PythonEnvironment } from '../../pythonEnvironments/info'; export const IEnvironmentActivationService = Symbol('IEnvironmentActivationService'); export interface IEnvironmentActivationService { + getProcessEnvironmentVariables(resource: Resource, shell?: string): Promise; getActivatedEnvironmentVariables( resource: Resource, interpreter?: PythonEnvironment, diff --git a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts index cb0b6b02f288..458e2b28c181 100644 --- a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts +++ b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts @@ -5,10 +5,10 @@ import * as sinon from 'sinon'; import { assert, expect } from 'chai'; -import { cloneDeep } from 'lodash'; import { mock, instance, when, anything, verify, reset } from 'ts-mockito'; import { EnvironmentVariableCollection, + EnvironmentVariableMutatorOptions, EnvironmentVariableScope, ProgressLocation, Uri, @@ -31,10 +31,7 @@ import { import { Interpreters } from '../../../client/common/utils/localize'; import { OSType, getOSType } from '../../../client/common/utils/platform'; import { defaultShells } from '../../../client/interpreter/activation/service'; -import { - TerminalEnvVarCollectionService, - _normCaseKeys, -} from '../../../client/interpreter/activation/terminalEnvVarCollectionService'; +import { TerminalEnvVarCollectionService } from '../../../client/interpreter/activation/terminalEnvVarCollectionService'; import { IEnvironmentActivationService } from '../../../client/interpreter/activation/types'; import { IInterpreterService } from '../../../client/interpreter/contracts'; import { PathUtils } from '../../../client/common/platform/pathUtils'; @@ -89,11 +86,16 @@ suite('Terminal Environment Variable Collection Service', () => { }) .thenResolve(); environmentActivationService = mock(); + when(environmentActivationService.getProcessEnvironmentVariables(anything(), anything())).thenResolve( + process.env, + ); configService = mock(); when(configService.getSettings(anything())).thenReturn(({ terminal: { activateEnvironment: true }, pythonPath: displayPath, } as unknown) as IPythonSettings); + when(collection.clear()).thenResolve(); + when(scopedCollection.clear()).thenResolve(); terminalEnvVarCollectionService = new TerminalEnvVarCollectionService( instance(platform), instance(interpreterService), @@ -169,8 +171,27 @@ suite('Terminal Environment Variable Collection Service', () => { }); test('If activated variables are returned for custom shell, apply it correctly to the collection', async () => { - const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ..._normCaseKeys(process.env) }; - delete envVars.PATH; + const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; + when( + environmentActivationService.getActivatedEnvironmentVariables( + anything(), + undefined, + undefined, + customShell, + ), + ).thenResolve(envVars); + + when(collection.replace(anything(), anything(), anything())).thenResolve(); + when(collection.delete(anything())).thenResolve(); + + await terminalEnvVarCollectionService._applyCollection(undefined, customShell); + + verify(collection.clear()).once(); + verify(collection.replace('CONDA_PREFIX', 'prefix/to/conda', anything())).once(); + }); + + test('If activated variables contain PS1, prefix it using shell integration', async () => { + const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env, PS1: '(prompt)' }; when( environmentActivationService.getActivatedEnvironmentVariables( anything(), @@ -182,16 +203,20 @@ suite('Terminal Environment Variable Collection Service', () => { when(collection.replace(anything(), anything(), anything())).thenResolve(); when(collection.delete(anything())).thenResolve(); + let opts: EnvironmentVariableMutatorOptions | undefined; + when(collection.prepend('PS1', '(prompt)', anything())).thenCall((_, _v, o) => { + opts = o; + }); await terminalEnvVarCollectionService._applyCollection(undefined, customShell); + verify(collection.clear()).once(); verify(collection.replace('CONDA_PREFIX', 'prefix/to/conda', anything())).once(); - verify(collection.delete('PATH')).once(); + assert.deepEqual(opts, { applyAtProcessCreation: false, applyAtShellIntegration: true }); }); test('Verify envs are not applied if env activation is disabled', async () => { - const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ..._normCaseKeys(process.env) }; - delete envVars.PATH; + const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; when( environmentActivationService.getActivatedEnvironmentVariables( anything(), @@ -211,13 +236,12 @@ suite('Terminal Environment Variable Collection Service', () => { await terminalEnvVarCollectionService._applyCollection(undefined, customShell); + verify(collection.clear()).once(); verify(collection.replace('CONDA_PREFIX', 'prefix/to/conda', anything())).never(); - verify(collection.delete('PATH')).never(); }); test('Verify correct options are used when applying envs and setting description', async () => { - const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ..._normCaseKeys(process.env) }; - delete envVars.PATH; + const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; const resource = Uri.file('a'); const workspaceFolder: WorkspaceFolder = { uri: Uri.file('workspacePath'), @@ -233,51 +257,11 @@ suite('Terminal Environment Variable Collection Service', () => { assert.deepEqual(options, { applyAtShellIntegration: true }); return Promise.resolve(); }); - when(scopedCollection.delete(anything())).thenResolve(); await terminalEnvVarCollectionService._applyCollection(resource, customShell); + verify(scopedCollection.clear()).once(); verify(scopedCollection.replace('CONDA_PREFIX', 'prefix/to/conda', anything())).once(); - verify(scopedCollection.delete('PATH')).once(); - }); - - test('Only relative changes to previously applied variables are applied to the collection', async () => { - const envVars: NodeJS.ProcessEnv = { - RANDOM_VAR: 'random', - CONDA_PREFIX: 'prefix/to/conda', - ..._normCaseKeys(process.env), - }; - when( - environmentActivationService.getActivatedEnvironmentVariables( - anything(), - undefined, - undefined, - customShell, - ), - ).thenResolve(envVars); - - when(collection.replace(anything(), anything(), anything())).thenResolve(); - when(collection.delete(anything())).thenResolve(); - - await terminalEnvVarCollectionService._applyCollection(undefined, customShell); - - const newEnvVars = cloneDeep(envVars); - delete newEnvVars.CONDA_PREFIX; - newEnvVars.RANDOM_VAR = undefined; // Deleting the variable from the collection is the same as setting it to undefined. - reset(environmentActivationService); - when( - environmentActivationService.getActivatedEnvironmentVariables( - anything(), - undefined, - undefined, - customShell, - ), - ).thenResolve(newEnvVars); - - await terminalEnvVarCollectionService._applyCollection(undefined, customShell); - - verify(collection.delete('CONDA_PREFIX')).once(); - verify(collection.delete('RANDOM_VAR')).once(); }); test('If no activated variables are returned for custom shell, fallback to using default shell', async () => { @@ -289,7 +273,7 @@ suite('Terminal Environment Variable Collection Service', () => { customShell, ), ).thenResolve(undefined); - const envVars = { CONDA_PREFIX: 'prefix/to/conda', ..._normCaseKeys(process.env) }; + const envVars = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; when( environmentActivationService.getActivatedEnvironmentVariables( anything(), @@ -300,12 +284,11 @@ suite('Terminal Environment Variable Collection Service', () => { ).thenResolve(envVars); when(collection.replace(anything(), anything(), anything())).thenResolve(); - when(collection.delete(anything())).thenResolve(); await terminalEnvVarCollectionService._applyCollection(undefined, customShell); verify(collection.replace('CONDA_PREFIX', 'prefix/to/conda', anything())).once(); - verify(collection.delete(anything())).never(); + verify(collection.clear()).twice(); }); test('If no activated variables are returned for default shell, clear collection', async () => {