Skip to content

Commit 6af3d03

Browse files
authored
Refresh environments immediately (#23634)
1 parent 8a87e1d commit 6af3d03

File tree

2 files changed

+138
-43
lines changed

2 files changed

+138
-43
lines changed

src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts

Lines changed: 137 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,24 @@
11
// Copyright (c) Microsoft Corporation. All rights reserved.
22
// Licensed under the MIT License.
33

4-
import { Disposable, EventEmitter, Event, Uri, workspace } from 'vscode';
4+
import { Disposable, EventEmitter, Event, workspace, Uri } from 'vscode';
55
import * as ch from 'child_process';
66
import * as path from 'path';
77
import * as rpc from 'vscode-jsonrpc/node';
88
import { PassThrough } from 'stream';
99
import { isWindows } from '../../../../common/platform/platformService';
1010
import { EXTENSION_ROOT_DIR } from '../../../../constants';
1111
import { traceError, traceInfo, traceLog, traceVerbose, traceWarn } from '../../../../logging';
12-
import { createDeferred } from '../../../../common/utils/async';
12+
import { createDeferred, createDeferredFrom } from '../../../../common/utils/async';
1313
import { DisposableBase, DisposableStore } from '../../../../common/utils/resourceLifecycle';
14-
import { getPythonSetting } from '../../../common/externalDependencies';
1514
import { DEFAULT_INTERPRETER_PATH_SETTING_KEY } from '../lowLevel/customWorkspaceLocator';
1615
import { noop } from '../../../../common/utils/misc';
16+
import { getConfiguration } from '../../../../common/vscodeApis/workspaceApis';
17+
import { CONDAPATH_SETTING_KEY } from '../../../common/environmentManagers/conda';
18+
import { VENVFOLDERS_SETTING_KEY, VENVPATH_SETTING_KEY } from '../lowLevel/customVirtualEnvLocator';
19+
import { getUserHomeDir } from '../../../../common/utils/platform';
20+
21+
const untildify = require('untildify');
1722

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

4449
export interface NativeGlobalPythonFinder extends Disposable {
4550
resolve(executable: string): Promise<NativeEnvInfo>;
46-
refresh(paths: Uri[]): AsyncIterable<NativeEnvInfo>;
51+
refresh(): AsyncIterable<NativeEnvInfo>;
4752
}
4853

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

62+
private firstRefreshResults: undefined | (() => AsyncGenerator<NativeEnvInfo, void, unknown>);
63+
5764
constructor() {
5865
super();
5966
this.connection = this.start();
67+
this.firstRefreshResults = this.refreshFirstTime();
6068
}
6169

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

74-
async *refresh(_paths: Uri[]): AsyncIterable<NativeEnvInfo> {
82+
async *refresh(): AsyncIterable<NativeEnvInfo> {
83+
if (this.firstRefreshResults) {
84+
// If this is the first time we are refreshing,
85+
// Then get the results from the first refresh.
86+
// Those would have started earlier and cached in memory.
87+
const results = this.firstRefreshResults();
88+
this.firstRefreshResults = undefined;
89+
yield* results;
90+
} else {
91+
const result = this.doRefresh();
92+
let completed = false;
93+
void result.completed.finally(() => {
94+
completed = true;
95+
});
96+
const envs: NativeEnvInfo[] = [];
97+
let discovered = createDeferred();
98+
const disposable = result.discovered((data) => {
99+
envs.push(data);
100+
discovered.resolve();
101+
});
102+
do {
103+
if (!envs.length) {
104+
await Promise.race([result.completed, discovered.promise]);
105+
}
106+
if (envs.length) {
107+
const dataToSend = [...envs];
108+
envs.length = 0;
109+
for (const data of dataToSend) {
110+
yield data;
111+
}
112+
}
113+
if (!completed) {
114+
discovered = createDeferred();
115+
}
116+
} while (!completed);
117+
disposable.dispose();
118+
}
119+
}
120+
121+
refreshFirstTime() {
75122
const result = this.doRefresh();
76-
let completed = false;
77-
void result.completed.finally(() => {
78-
completed = true;
79-
});
123+
const completed = createDeferredFrom(result.completed);
80124
const envs: NativeEnvInfo[] = [];
81125
let discovered = createDeferred();
82126
const disposable = result.discovered((data) => {
83127
envs.push(data);
84128
discovered.resolve();
85129
});
86-
do {
87-
if (!envs.length) {
88-
await Promise.race([result.completed, discovered.promise]);
89-
}
90-
if (envs.length) {
91-
const dataToSend = [...envs];
92-
envs.length = 0;
93-
for (const data of dataToSend) {
94-
yield data;
130+
131+
const iterable = async function* () {
132+
do {
133+
if (!envs.length) {
134+
await Promise.race([completed.promise, discovered.promise]);
135+
}
136+
if (envs.length) {
137+
const dataToSend = [...envs];
138+
envs.length = 0;
139+
for (const data of dataToSend) {
140+
yield data;
141+
}
95142
}
96-
}
97-
if (!completed) {
98-
discovered = createDeferred();
99-
}
100-
} while (!completed);
101-
disposable.dispose();
143+
if (!completed.completed) {
144+
discovered = createDeferred();
145+
}
146+
} while (!completed.completed);
147+
disposable.dispose();
148+
};
149+
150+
return iterable.bind(this);
102151
}
103152

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

224-
const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
225-
getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri.fsPath),
226-
);
227-
pythonPathSettings.push(getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
228-
const pythonSettings = Array.from(new Set(pythonPathSettings.filter((item) => !!item)).values()).map((p) =>
229-
// We only want the parent directories.
230-
path.dirname(p!),
231-
);
232273
trackPromiseAndNotifyOnCompletion(
233-
this.connection
234-
.sendRequest<{ duration: number }>('refresh', {
235-
// Send configuration information to the Python finder.
236-
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
237-
// Also send the python paths that are configured in the settings.
238-
python_path_settings: pythonSettings,
239-
conda_executable: undefined,
240-
})
274+
this.sendRefreshRequest()
241275
.then(({ duration }) => traceInfo(`Native Python Finder completed in ${duration}ms`))
242276
.catch((ex) => traceError('Error in Native Python Finder', ex)),
243277
);
278+
244279
completed.promise.finally(() => disposable.dispose());
245280
return {
246281
completed: completed.promise,
247282
discovered: discovered.event,
248283
};
249284
}
285+
286+
private sendRefreshRequest() {
287+
const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
288+
getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri),
289+
);
290+
pythonPathSettings.push(getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
291+
// We can have multiple workspaces, each with its own setting.
292+
const pythonSettings = Array.from(
293+
new Set(
294+
pythonPathSettings
295+
.filter((item) => !!item)
296+
// We only want the parent directories.
297+
.map((p) => path.dirname(p!))
298+
/// If setting value is 'python', then `path.dirname('python')` will yield `.`
299+
.filter((item) => item !== '.'),
300+
),
301+
);
302+
303+
return this.connection.sendRequest<{ duration: number }>(
304+
'refresh',
305+
// Send configuration information to the Python finder.
306+
{
307+
// This has a special meaning in locator, its lot a low priority
308+
// as we treat this as workspace folders that can contain a large number of files.
309+
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
310+
// Also send the python paths that are configured in the settings.
311+
python_interpreter_paths: pythonSettings,
312+
// We do not want to mix this with `search_paths`
313+
virtual_env_paths: getCustomVirtualEnvDirs(),
314+
conda_executable: getPythonSettingAndUntildify<string>(CONDAPATH_SETTING_KEY),
315+
poetry_executable: getPythonSettingAndUntildify<string>('poetryPath'),
316+
pipenv_executable: getPythonSettingAndUntildify<string>('pipenvPath'),
317+
},
318+
);
319+
}
320+
}
321+
322+
/**
323+
* Gets all custom virtual environment locations to look for environments.
324+
*/
325+
async function getCustomVirtualEnvDirs(): Promise<string[]> {
326+
const venvDirs: string[] = [];
327+
const venvPath = getPythonSettingAndUntildify<string>(VENVPATH_SETTING_KEY);
328+
if (venvPath) {
329+
venvDirs.push(untildify(venvPath));
330+
}
331+
const venvFolders = getPythonSettingAndUntildify<string[]>(VENVFOLDERS_SETTING_KEY) ?? [];
332+
const homeDir = getUserHomeDir();
333+
if (homeDir) {
334+
venvFolders.map((item) => path.join(homeDir, item)).forEach((d) => venvDirs.push(d));
335+
}
336+
return Array.from(new Set(venvDirs));
337+
}
338+
339+
function getPythonSettingAndUntildify<T>(name: string, scope?: Uri): T | undefined {
340+
const value = getConfiguration('python', scope).get<T>(name);
341+
if (typeof value === 'string') {
342+
return value ? ((untildify(value as string) as unknown) as T) : undefined;
343+
}
344+
return value;
250345
}
251346

252347
export function createNativeGlobalPythonFinder(): NativeGlobalPythonFinder {

src/client/pythonEnvironments/base/locators/lowLevel/nativeLocator.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ export class NativeLocator implements ILocator<BasicEnvInfo>, IDisposable {
105105
const disposables: IDisposable[] = [];
106106
const disposable = new Disposable(() => disposeAll(disposables));
107107
this.disposables.push(disposable);
108-
for await (const data of this.finder.refresh([])) {
108+
for await (const data of this.finder.refresh()) {
109109
if (data.manager) {
110110
switch (toolToKnownEnvironmentTool(data.manager.tool)) {
111111
case 'Conda': {

0 commit comments

Comments
 (0)