Improvements to pythonTerminalEnvVarActivation experiment (#21751)

Kartik Raj · web-flow · commit 23353bbd1bbe · 2023-08-04T08:52:05.000-07:00
diff --git 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<NodeJS.ProcessEnv | undefined>(CACHE_DURATION);
+        const memCache = new InMemoryCache<NodeJS.ProcessEnv | undefined>(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<EnvironmentVariables> {
+        // 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
@@ -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<void> {
         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
@@ -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<EnvironmentVariables>;
     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
