Skip to content

Commit

Permalink
Refresh environments immediately (microsoft#23634)
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne authored Jun 18, 2024
1 parent 8a87e1d commit 6af3d03
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 43 deletions.
179 changes: 137 additions & 42 deletions src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,24 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import { Disposable, EventEmitter, Event, Uri, workspace } from 'vscode';
import { Disposable, EventEmitter, Event, workspace, Uri } from 'vscode';
import * as ch from 'child_process';
import * as path from 'path';
import * as rpc from 'vscode-jsonrpc/node';
import { PassThrough } from 'stream';
import { isWindows } from '../../../../common/platform/platformService';
import { EXTENSION_ROOT_DIR } from '../../../../constants';
import { traceError, traceInfo, traceLog, traceVerbose, traceWarn } from '../../../../logging';
import { createDeferred } from '../../../../common/utils/async';
import { createDeferred, createDeferredFrom } from '../../../../common/utils/async';
import { DisposableBase, DisposableStore } from '../../../../common/utils/resourceLifecycle';
import { getPythonSetting } from '../../../common/externalDependencies';
import { DEFAULT_INTERPRETER_PATH_SETTING_KEY } from '../lowLevel/customWorkspaceLocator';
import { noop } from '../../../../common/utils/misc';
import { getConfiguration } from '../../../../common/vscodeApis/workspaceApis';
import { CONDAPATH_SETTING_KEY } from '../../../common/environmentManagers/conda';
import { VENVFOLDERS_SETTING_KEY, VENVPATH_SETTING_KEY } from '../lowLevel/customVirtualEnvLocator';
import { getUserHomeDir } from '../../../../common/utils/platform';

const untildify = require('untildify');

const NATIVE_LOCATOR = isWindows()
? path.join(EXTENSION_ROOT_DIR, 'native_locator', 'bin', 'pet.exe')
Expand Down Expand Up @@ -43,7 +48,7 @@ export interface NativeEnvManagerInfo {

export interface NativeGlobalPythonFinder extends Disposable {
resolve(executable: string): Promise<NativeEnvInfo>;
refresh(paths: Uri[]): AsyncIterable<NativeEnvInfo>;
refresh(): AsyncIterable<NativeEnvInfo>;
}

interface NativeLog {
Expand All @@ -54,9 +59,12 @@ interface NativeLog {
class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGlobalPythonFinder {
private readonly connection: rpc.MessageConnection;

private firstRefreshResults: undefined | (() => AsyncGenerator<NativeEnvInfo, void, unknown>);

constructor() {
super();
this.connection = this.start();
this.firstRefreshResults = this.refreshFirstTime();
}

public async resolve(executable: string): Promise<NativeEnvInfo> {
Expand All @@ -71,41 +79,82 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
return environment;
}

async *refresh(_paths: Uri[]): AsyncIterable<NativeEnvInfo> {
async *refresh(): AsyncIterable<NativeEnvInfo> {
if (this.firstRefreshResults) {
// If this is the first time we are refreshing,
// Then get the results from the first refresh.
// Those would have started earlier and cached in memory.
const results = this.firstRefreshResults();
this.firstRefreshResults = undefined;
yield* results;
} else {
const result = this.doRefresh();
let completed = false;
void result.completed.finally(() => {
completed = true;
});
const envs: NativeEnvInfo[] = [];
let discovered = createDeferred();
const disposable = result.discovered((data) => {
envs.push(data);
discovered.resolve();
});
do {
if (!envs.length) {
await Promise.race([result.completed, discovered.promise]);
}
if (envs.length) {
const dataToSend = [...envs];
envs.length = 0;
for (const data of dataToSend) {
yield data;
}
}
if (!completed) {
discovered = createDeferred();
}
} while (!completed);
disposable.dispose();
}
}

refreshFirstTime() {
const result = this.doRefresh();
let completed = false;
void result.completed.finally(() => {
completed = true;
});
const completed = createDeferredFrom(result.completed);
const envs: NativeEnvInfo[] = [];
let discovered = createDeferred();
const disposable = result.discovered((data) => {
envs.push(data);
discovered.resolve();
});
do {
if (!envs.length) {
await Promise.race([result.completed, discovered.promise]);
}
if (envs.length) {
const dataToSend = [...envs];
envs.length = 0;
for (const data of dataToSend) {
yield data;

const iterable = async function* () {
do {
if (!envs.length) {
await Promise.race([completed.promise, discovered.promise]);
}
if (envs.length) {
const dataToSend = [...envs];
envs.length = 0;
for (const data of dataToSend) {
yield data;
}
}
}
if (!completed) {
discovered = createDeferred();
}
} while (!completed);
disposable.dispose();
if (!completed.completed) {
discovered = createDeferred();
}
} while (!completed.completed);
disposable.dispose();
};

return iterable.bind(this);
}

// eslint-disable-next-line class-methods-use-this
private start(): rpc.MessageConnection {
const proc = ch.spawn(NATIVE_LOCATOR, ['server'], { env: process.env });
const disposables: Disposable[] = [];
// jsonrpc package cannot handle messages coming through too quicly.
// jsonrpc package cannot handle messages coming through too quickly.
// Lets handle the messages and close the stream only when
// we have got the exit event.
const readable = new PassThrough();
Expand Down Expand Up @@ -213,40 +262,86 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
discovered.fire(environment);
})
.catch((ex) => traceError(`Error in Resolving Python Environment ${data}`, ex));
.catch((ex) => traceError(`Error in Resolving Python Environment ${JSON.stringify(data)}`, ex));
trackPromiseAndNotifyOnCompletion(promise);
} else {
discovered.fire(data);
}
}),
);

const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri.fsPath),
);
pythonPathSettings.push(getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
const pythonSettings = Array.from(new Set(pythonPathSettings.filter((item) => !!item)).values()).map((p) =>
// We only want the parent directories.
path.dirname(p!),
);
trackPromiseAndNotifyOnCompletion(
this.connection
.sendRequest<{ duration: number }>('refresh', {
// Send configuration information to the Python finder.
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
// Also send the python paths that are configured in the settings.
python_path_settings: pythonSettings,
conda_executable: undefined,
})
this.sendRefreshRequest()
.then(({ duration }) => traceInfo(`Native Python Finder completed in ${duration}ms`))
.catch((ex) => traceError('Error in Native Python Finder', ex)),
);

completed.promise.finally(() => disposable.dispose());
return {
completed: completed.promise,
discovered: discovered.event,
};
}

private sendRefreshRequest() {
const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri),
);
pythonPathSettings.push(getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
// We can have multiple workspaces, each with its own setting.
const pythonSettings = Array.from(
new Set(
pythonPathSettings
.filter((item) => !!item)
// We only want the parent directories.
.map((p) => path.dirname(p!))
/// If setting value is 'python', then `path.dirname('python')` will yield `.`
.filter((item) => item !== '.'),
),
);

return this.connection.sendRequest<{ duration: number }>(
'refresh',
// Send configuration information to the Python finder.
{
// This has a special meaning in locator, its lot a low priority
// as we treat this as workspace folders that can contain a large number of files.
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
// Also send the python paths that are configured in the settings.
python_interpreter_paths: pythonSettings,
// We do not want to mix this with `search_paths`
virtual_env_paths: getCustomVirtualEnvDirs(),
conda_executable: getPythonSettingAndUntildify<string>(CONDAPATH_SETTING_KEY),
poetry_executable: getPythonSettingAndUntildify<string>('poetryPath'),
pipenv_executable: getPythonSettingAndUntildify<string>('pipenvPath'),
},
);
}
}

/**
* Gets all custom virtual environment locations to look for environments.
*/
async function getCustomVirtualEnvDirs(): Promise<string[]> {
const venvDirs: string[] = [];
const venvPath = getPythonSettingAndUntildify<string>(VENVPATH_SETTING_KEY);
if (venvPath) {
venvDirs.push(untildify(venvPath));
}
const venvFolders = getPythonSettingAndUntildify<string[]>(VENVFOLDERS_SETTING_KEY) ?? [];
const homeDir = getUserHomeDir();
if (homeDir) {
venvFolders.map((item) => path.join(homeDir, item)).forEach((d) => venvDirs.push(d));
}
return Array.from(new Set(venvDirs));
}

function getPythonSettingAndUntildify<T>(name: string, scope?: Uri): T | undefined {
const value = getConfiguration('python', scope).get<T>(name);
if (typeof value === 'string') {
return value ? ((untildify(value as string) as unknown) as T) : undefined;
}
return value;
}

export function createNativeGlobalPythonFinder(): NativeGlobalPythonFinder {
Expand Down
2 changes: 1 addition & 1 deletion src/client/pythonEnvironments/base/locators/lowLevel/nativeLocator.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ export class NativeLocator implements ILocator<BasicEnvInfo>, IDisposable {
const disposables: IDisposable[] = [];
const disposable = new Disposable(() => disposeAll(disposables));
this.disposables.push(disposable);
for await (const data of this.finder.refresh([])) {
for await (const data of this.finder.refresh()) {
if (data.manager) {
switch (toolToKnownEnvironmentTool(data.manager.tool)) {
case 'Conda': {
Expand Down

0 comments on commit 6af3d03

Please sign in to comment.