Skip to content

Commit

Permalink
Use python.envFile to get env file name when running Python code (#10876
Browse files Browse the repository at this point in the history
)

* Use python.envFile for python kernels

* Added tests

* Misc

* Misc

* Address code review comments

* Logging

* oops

* Misc

* Misc

* Misc

* Misc

* Fixes

* Misc

* oops

* Misc

* Misc

* Misc

* oops

* Misc

* misc

* Misc
  • Loading branch information
DonJayamanne authored Jul 21, 2022
1 parent 728cc75 commit 7b12956
Show file tree
Hide file tree
Showing 15 changed files with 474 additions and 85 deletions.
2 changes: 1 addition & 1 deletion src/kernels/raw/finder/jupyterPaths.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ export class JupyterPaths {
subDir?: string
): Promise<Uri[]> {
const paths = new ResourceSet();
const vars = await this.envVarsProvider.getEnvironmentVariables();
const vars = await this.envVarsProvider.getEnvironmentVariables(undefined, 'RunPythonCode');
if (cancelToken?.isCancellationRequested) {
return [];
}
Expand Down
7 changes: 5 additions & 2 deletions src/kernels/raw/launcher/kernelEnvVarsService.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import { IInterpreterService } from '../../../platform/interpreter/contracts';
import { PythonEnvironment } from '../../../platform/pythonEnvironments/info';
import { IJupyterKernelSpec } from '../../types';
import { Uri } from 'vscode';
import { PYTHON_LANGUAGE } from '../../../platform/common/constants';

/**
* Class used to fetch environment variables for a kernel.
Expand Down Expand Up @@ -49,7 +50,7 @@ export class KernelEnvironmentVariablesService {
kernelSpec: IJupyterKernelSpec
) {
let kernelEnv = kernelSpec.env && Object.keys(kernelSpec.env).length > 0 ? kernelSpec.env : undefined;

const isPythonKernel = (kernelSpec.language || '').toLowerCase() === PYTHON_LANGUAGE;
// If an interpreter was not explicitly passed in, check for an interpreter path in the kernelspec to use
if (!interpreter && kernelSpec.interpreterPath) {
interpreter = await this.interpreterService
Expand All @@ -61,7 +62,9 @@ export class KernelEnvironmentVariablesService {
}

let [customEnvVars, interpreterEnv] = await Promise.all([
this.customEnvVars.getCustomEnvironmentVariables(resource).catch(noop),
this.customEnvVars
.getCustomEnvironmentVariables(resource, isPythonKernel ? 'RunPythonCode' : 'RunNonPythonCode')
.catch(noop),
interpreter
? this.envActivation
.getActivatedEnvironmentVariables(resource, interpreter, false)
Expand Down
6 changes: 4 additions & 2 deletions src/kernels/variables/preWarmVariables.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,11 @@ export class PreWarmActivatedJupyterEnvironmentVariables implements IExtensionSi
// Don't try to pre-warm variables if user has too many workspace folders opened.
const workspaceFolderCount = this.workspace.workspaceFolders?.length ?? 0;
if (workspaceFolderCount <= 5) {
this.envVarsProvider.getEnvironmentVariables(undefined).ignoreErrors();
this.envVarsProvider.getEnvironmentVariables(undefined, 'RunNonPythonCode').ignoreErrors();
this.envVarsProvider.getEnvironmentVariables(undefined, 'RunPythonCode').ignoreErrors();
(this.workspace.workspaceFolders || []).forEach((folder) => {
this.envVarsProvider.getEnvironmentVariables(folder.uri).ignoreErrors();
this.envVarsProvider.getEnvironmentVariables(folder.uri, 'RunNonPythonCode').ignoreErrors();
this.envVarsProvider.getEnvironmentVariables(folder.uri, 'RunPythonCode').ignoreErrors();
});
}
this.condaService.getCondaFile().ignoreErrors();
Expand Down
40 changes: 17 additions & 23 deletions src/platform/common/process/environmentActivationService.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -130,19 +130,9 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
) {
EnvironmentActivationService.minTimeAfterWhichWeShouldCacheEnvVariables =
minTimeAfterWhichWeShouldCacheEnvVariables;
this.envVarsService.onDidEnvironmentVariablesChange(
() => this.activatedEnvVariablesCache.clear(),
this,
this.disposables
);

this.interpreterService.onDidChangeInterpreter(
() => this.activatedEnvVariablesCache.clear(),
this,
this.disposables
);
this.envVarsService.onDidEnvironmentVariablesChange(this.clearCache, this, this.disposables);
this.interpreterService.onDidChangeInterpreter(this.clearCache, this, this.disposables);
}
@testOnlyMethod()
public clearCache() {
this.activatedEnvVariablesCache.clear();
}
Expand Down Expand Up @@ -195,6 +185,11 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
envVariablesFromPython.completed &&
!envVariablesFromPython.value
) {
traceWarning(
`Failed to get env vars from Python extension. Falling back to ours for ${getDisplayPath(
interpreter.uri
)}.`
);
await envVariablesOurSelves.promise;
}

Expand All @@ -206,8 +201,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
} else {
traceVerbose(`Got env vars ourselves faster, but empty ${getDisplayPath(interpreter?.uri)}`);
}
}
if (!envVariablesOurSelves.resolved) {
} else {
traceVerbose(`Got env vars with python ext faster ${getDisplayPath(interpreter?.uri)}`);
}
return envVariablesFromPython.promise;
Expand All @@ -222,7 +216,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
): Promise<NodeJS.ProcessEnv | undefined> {
const stopWatch = new StopWatch();
// We'll need this later.
this.envVarsService.getEnvironmentVariables(resource).catch(noop);
this.envVarsService.getEnvironmentVariables(resource, 'RunPythonCode').catch(noop);

// Check cache.
let reasonForFailure:
Expand All @@ -244,7 +238,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
return undefined;
})
),
this.envVarsService.getCustomEnvironmentVariables(resource).catch((ex) => {
this.envVarsService.getCustomEnvironmentVariables(resource, 'RunPythonCode').catch((ex) => {
traceError(
`Failed to get activated env variables from Python Extension for ${getDisplayPath(
interpreter.uri
Expand Down Expand Up @@ -301,6 +295,11 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
const workspaceKey = this.workspace.getWorkspaceFolderIdentifier(resource);
const key = `${workspaceKey}_${interpreter && getInterpreterHash(interpreter)}`;

if (this.activatedEnvVariablesCache.has(key)) {
traceVerbose(`Got activation Env Vars from cached promise with key ${key}`);
return this.activatedEnvVariablesCache.get(key);
}

const shellInfo = defaultShells[this.platform.osType];
const envType = interpreter?.envType;
if (!shellInfo) {
Expand All @@ -319,11 +318,6 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
return;
}

if (this.activatedEnvVariablesCache.has(key)) {
traceVerbose(`Got activation Env Vars from cached promise with key ${key}`);
return this.activatedEnvVariablesCache.get(key);
}

const promise = (async () => {
const condaActivation = async () => {
const stopWatch = new StopWatch();
Expand Down Expand Up @@ -378,7 +372,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi

const [activationCommands, customEnvVars] = await Promise.all([
this.getActivationCommands(resource, interpreterDetails || interpreter),
this.envVarsService.getEnvironmentVariables(resource)
this.envVarsService.getEnvironmentVariables(resource, 'RunPythonCode')
]);
const processService = await processServicePromise;
const hasCustomEnvVars = Object.keys(customEnvVars).length;
Expand Down Expand Up @@ -573,7 +567,7 @@ export class EnvironmentActivationService implements IEnvironmentActivationServi
this.condaService.getCondaFile(),
this.condaService.getCondaVersion(),
this.fs.createTemporaryLocalFile('.json'),
this.envVarsService.getEnvironmentVariables(resource)
this.envVarsService.getEnvironmentVariables(resource, 'RunPythonCode')
]);
const hasCustomEnvVars = Object.keys(customEnvVars).length;
const env = hasCustomEnvVars
Expand Down
2 changes: 1 addition & 1 deletion src/platform/common/process/processFactory.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ export class ProcessServiceFactory implements IProcessServiceFactory {
if (!this.workspace.isTrusted) {
throw new Error('Workspace not trusted');
}
const customEnvVars = await this.envVarsService.getEnvironmentVariables(resource);
const customEnvVars = await this.envVarsService.getEnvironmentVariables(resource, 'RunNonPythonCode');
const proc: IProcessService = new ProcessService(this.decoder, customEnvVars);
this.disposableRegistry.push(proc);
return proc.on('exec', this.processLogger.logProcess.bind(this.processLogger));
Expand Down
136 changes: 101 additions & 35 deletions src/platform/common/variables/customEnvironmentVariablesProvider.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,17 @@
// Licensed under the MIT License.
import { inject, injectable, optional } from 'inversify';
import * as path from '../../vscode-path/path';
import { ConfigurationChangeEvent, Disposable, Event, EventEmitter, FileSystemWatcher, Uri } from 'vscode';
import { ConfigurationChangeEvent, Disposable, Event, EventEmitter, RelativePattern, Uri } from 'vscode';
import { TraceOptions } from '../../logging/types';
import { sendFileCreationTelemetry } from '../../telemetry/envFileTelemetry.node';
import { IWorkspaceService } from '../application/types';
import { IDisposableRegistry } from '../types';
import { IDisposableRegistry, Resource } from '../types';
import { InMemoryCache } from '../utils/cacheUtils';
import { EnvironmentVariables, ICustomEnvironmentVariablesProvider, IEnvironmentVariablesService } from './types';
import { traceDecoratorVerbose, traceInfoIfCI, traceVerbose } from '../../logging';
import { traceDecoratorVerbose, traceError, traceInfoIfCI, traceVerbose } from '../../logging';
import { SystemVariables } from './systemVariables.node';
import { disposeAllDisposables } from '../helpers';
import { IPythonExtensionChecker } from '../../api/types';

const CACHE_DURATION = 60 * 60 * 1000;
@injectable()
Expand All @@ -18,58 +21,76 @@ export class CustomEnvironmentVariablesProvider implements ICustomEnvironmentVar
return this.changeEventEmitter.event;
}
public trackedWorkspaceFolders = new Set<string>();
private fileWatchers = new Map<string, FileSystemWatcher>();
private fileWatchers = new Set<string>();
private disposables: Disposable[] = [];
private changeEventEmitter: EventEmitter<Uri | undefined>;
private changeEventEmitter = new EventEmitter<Uri | undefined>();
constructor(
@inject(IEnvironmentVariablesService) private envVarsService: IEnvironmentVariablesService,
@inject(IDisposableRegistry) disposableRegistry: Disposable[],
@inject(IWorkspaceService) private workspaceService: IWorkspaceService,
@inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker,
@inject('number') @optional() private cacheDuration: number = CACHE_DURATION
) {
disposableRegistry.push(this);
this.changeEventEmitter = new EventEmitter();
const disposable = this.workspaceService.onDidChangeConfiguration(this.configurationChanged, this);
this.disposables.push(disposable);
this.workspaceService.onDidChangeConfiguration(this.configurationChanged, this, this.disposables);
}

public dispose() {
this.changeEventEmitter.dispose();
this.fileWatchers.forEach((watcher) => {
if (watcher) {
watcher.dispose();
}
});
disposeAllDisposables(this.disposables);
}

@traceDecoratorVerbose('Get Custom Env Variables', TraceOptions.BeforeCall | TraceOptions.Arguments)
public async getEnvironmentVariables(resource?: Uri): Promise<EnvironmentVariables> {
public async getEnvironmentVariables(
resource: Resource,
purpose: 'RunPythonCode' | 'RunNonPythonCode'
): Promise<EnvironmentVariables> {
// Cache resource specific interpreter data
const key = this.workspaceService.getWorkspaceFolderIdentifier(resource);
const cacheStoreIndexedByWorkspaceFolder = new InMemoryCache<EnvironmentVariables>(this.cacheDuration, key);
const cacheStoreIndexedByWorkspaceFolder = new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForMergedVars(resource, purpose)
);

if (cacheStoreIndexedByWorkspaceFolder.hasData && cacheStoreIndexedByWorkspaceFolder.data) {
traceVerbose(`Cached data exists getEnvironmentVariables, ${resource ? resource.fsPath : '<No Resource>'}`);
return cacheStoreIndexedByWorkspaceFolder.data!;
}
const promise = this._getEnvironmentVariables(resource);
const promise = this._getEnvironmentVariables(resource, purpose);
promise.then((result) => (cacheStoreIndexedByWorkspaceFolder.data = result)).ignoreErrors();
return promise;
}
public async getCustomEnvironmentVariables(resource?: Uri): Promise<EnvironmentVariables | undefined> {
public async getCustomEnvironmentVariables(
resource: Resource,
purpose: 'RunPythonCode' | 'RunNonPythonCode'
): Promise<EnvironmentVariables | undefined> {
// Cache resource specific interpreter data
const cacheStoreIndexedByWorkspaceFolder = new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForCustomVars(resource, purpose)
);

if (cacheStoreIndexedByWorkspaceFolder.hasData) {
traceVerbose(
`Cached custom vars data exists getCustomEnvironmentVariables, ${
resource ? resource.fsPath : '<No Resource>'
}`
);
return cacheStoreIndexedByWorkspaceFolder.data!;
}

const workspaceFolderUri = this.getWorkspaceFolderUri(resource);
if (!workspaceFolderUri) {
traceInfoIfCI(`No workspace folder found for ${resource ? resource.fsPath : '<No Resource>'}`);
return;
}
this.trackedWorkspaceFolders.add(workspaceFolderUri ? workspaceFolderUri.fsPath : '');
this.trackedWorkspaceFolders.add(workspaceFolderUri.fsPath || '');

// eslint-disable-next-line
// TODO: This should be added to the python API (or this entire service should move there)
// https://github.com/microsoft/vscode-jupyter/issues/51
const envFile = path.join(workspaceFolderUri.fsPath, '.env');
const envFile = this.getEnvFile(workspaceFolderUri, purpose);
this.createFileWatcher(envFile, workspaceFolderUri);
return this.envVarsService.parseFile(envFile, process.env);

const promise = this.envVarsService.parseFile(envFile, process.env);
promise.then((result) => (cacheStoreIndexedByWorkspaceFolder.data = result)).ignoreErrors();
return promise;
}
public configurationChanged(e: ConfigurationChangeEvent) {
this.trackedWorkspaceFolders.forEach((item) => {
Expand All @@ -80,19 +101,43 @@ export class CustomEnvironmentVariablesProvider implements ICustomEnvironmentVar
});
}
public createFileWatcher(envFile: string, workspaceFolderUri?: Uri) {
if (this.fileWatchers.has(envFile)) {
const key = this.getCacheKeyForMergedVars(workspaceFolderUri, 'RunNonPythonCode');
if (this.fileWatchers.has(key)) {
return;
}
const envFileWatcher = this.workspaceService.createFileSystemWatcher(envFile);
this.fileWatchers.set(envFile, envFileWatcher);
const pattern = new RelativePattern(Uri.file(path.dirname(envFile)), path.basename(envFile));
const envFileWatcher = this.workspaceService.createFileSystemWatcher(pattern, false, false, false);
if (envFileWatcher) {
this.disposables.push(envFileWatcher.onDidChange(() => this.onEnvironmentFileChanged(workspaceFolderUri)));
this.disposables.push(envFileWatcher.onDidCreate(() => this.onEnvironmentFileCreated(workspaceFolderUri)));
this.disposables.push(envFileWatcher.onDidDelete(() => this.onEnvironmentFileChanged(workspaceFolderUri)));
this.disposables.push(envFileWatcher);
this.fileWatchers.add(key);
envFileWatcher.onDidChange(() => this.onEnvironmentFileChanged(workspaceFolderUri), this, this.disposables);
envFileWatcher.onDidCreate(() => this.onEnvironmentFileCreated(workspaceFolderUri), this, this.disposables);
envFileWatcher.onDidDelete(() => this.onEnvironmentFileChanged(workspaceFolderUri), this, this.disposables);
} else {
traceError('Failed to create file watcher for environment file');
}
}
private async _getEnvironmentVariables(resource?: Uri): Promise<EnvironmentVariables> {
let customEnvVars = await this.getCustomEnvironmentVariables(resource);
private getEnvFile(workspaceFolderUri: Uri, purpose: 'RunPythonCode' | 'RunNonPythonCode') {
if (purpose === 'RunPythonCode') {
return this.getPythonEnvFile(workspaceFolderUri) || path.join(workspaceFolderUri.fsPath, '.env');
} else {
return path.join(workspaceFolderUri.fsPath, '.env');
}
}
private getPythonEnvFile(resource?: Uri): string | undefined {
if (!this.extensionChecker.isPythonExtensionInstalled) {
return;
}
const workspaceFolderUri = this.getWorkspaceFolderUri(resource);
const sysVars = new SystemVariables(resource, workspaceFolderUri, this.workspaceService);
const pythonEnvSetting = this.workspaceService.getConfiguration('python', resource).get<string>('envFile');
return pythonEnvSetting ? sysVars.resolve(pythonEnvSetting) : undefined;
}
private async _getEnvironmentVariables(
resource: Resource,
purpose: 'RunPythonCode' | 'RunNonPythonCode'
): Promise<EnvironmentVariables> {
let customEnvVars = await this.getCustomEnvironmentVariables(resource, purpose);
if (!customEnvVars) {
customEnvVars = {};
}
Expand Down Expand Up @@ -127,10 +172,31 @@ export class CustomEnvironmentVariablesProvider implements ICustomEnvironmentVar
sendFileCreationTelemetry();
}

private onEnvironmentFileChanged(workspaceFolderUri?: Uri) {
const key = this.workspaceService.getWorkspaceFolderIdentifier(workspaceFolderUri);
new InMemoryCache<EnvironmentVariables>(this.cacheDuration, `$getEnvironmentVariables-${key}`).clear();
private onEnvironmentFileChanged(workspaceFolderUri: Resource) {
new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForMergedVars(workspaceFolderUri, 'RunNonPythonCode')
).clear();
new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForMergedVars(workspaceFolderUri, 'RunPythonCode')
).clear();

new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForCustomVars(workspaceFolderUri, 'RunNonPythonCode')
).clear();
new InMemoryCache<EnvironmentVariables>(
this.cacheDuration,
this.getCacheKeyForCustomVars(workspaceFolderUri, 'RunPythonCode')
).clear();

this.changeEventEmitter.fire(workspaceFolderUri);
}
private getCacheKeyForMergedVars(workspaceFolderUri: Resource, purpose: 'RunPythonCode' | 'RunNonPythonCode') {
return `${this.workspaceService.getWorkspaceFolderIdentifier(workspaceFolderUri)}:${purpose || ''}`;
}
private getCacheKeyForCustomVars(workspaceFolderUri: Resource, purpose: 'RunPythonCode' | 'RunNonPythonCode') {
return `${this.workspaceService.getWorkspaceFolderIdentifier(workspaceFolderUri)}:${purpose || ''}:CustomVars`;
}
}
Loading

0 comments on commit 7b12956

Please sign in to comment.