Skip to content

Commit

Permalink
Improvements to pythonTerminalEnvVarActivation experiment (microsof…
Browse files Browse the repository at this point in the history 
…t#21751)
  • Loading branch information
Kartik Raj authored Aug 4, 2023
1 parent 40bb62a commit 23353bb
Show file tree
Hide file tree
Showing 4 changed files with 118 additions and 90 deletions.
51 changes: 44 additions & 7 deletions src/client/interpreter/activation/service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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]}`;
Expand All @@ -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
Expand Down
58 changes: 32 additions & 26 deletions src/client/interpreter/activation/terminalEnvVarCollectionService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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,
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand All @@ -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];
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
2 changes: 2 additions & 0 deletions src/client/interpreter/activation/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading

0 comments on commit 23353bb

Please sign in to comment.