Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improvements to pythonTerminalEnvVarActivation experiment #21751

Merged
merged 16 commits into from
Aug 4, 2023
51 changes: 44 additions & 7 deletions src/client/interpreter/activation/service.ts
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
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
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