Skip to content

Commit 8a87e1d

Browse files
authored
Adopt native locator resolve methods (#23631)
1 parent 0835c7e commit 8a87e1d

File tree

3 files changed

+140
-48
lines changed

3 files changed

+140
-48
lines changed

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

Lines changed: 114 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@ import { isWindows } from '../../../../common/platform/platformService';
1010
import { EXTENSION_ROOT_DIR } from '../../../../constants';
1111
import { traceError, traceInfo, traceLog, traceVerbose, traceWarn } from '../../../../logging';
1212
import { createDeferred } from '../../../../common/utils/async';
13-
import { DisposableBase } from '../../../../common/utils/resourceLifecycle';
13+
import { DisposableBase, DisposableStore } from '../../../../common/utils/resourceLifecycle';
14+
import { getPythonSetting } from '../../../common/externalDependencies';
15+
import { DEFAULT_INTERPRETER_PATH_SETTING_KEY } from '../lowLevel/customWorkspaceLocator';
16+
import { noop } from '../../../../common/utils/misc';
1417

1518
const NATIVE_LOCATOR = isWindows()
1619
? path.join(EXTENSION_ROOT_DIR, 'native_locator', 'bin', 'pet.exe')
@@ -39,6 +42,7 @@ export interface NativeEnvManagerInfo {
3942
}
4043

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

@@ -48,18 +52,41 @@ interface NativeLog {
4852
}
4953

5054
class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGlobalPythonFinder {
55+
private readonly connection: rpc.MessageConnection;
56+
57+
constructor() {
58+
super();
59+
this.connection = this.start();
60+
}
61+
62+
public async resolve(executable: string): Promise<NativeEnvInfo> {
63+
const { environment, duration } = await this.connection.sendRequest<{
64+
duration: number;
65+
environment: NativeEnvInfo;
66+
}>('resolve', {
67+
executable,
68+
});
69+
70+
traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
71+
return environment;
72+
}
73+
5174
async *refresh(_paths: Uri[]): AsyncIterable<NativeEnvInfo> {
52-
const result = this.start();
75+
const result = this.doRefresh();
5376
let completed = false;
5477
void result.completed.finally(() => {
5578
completed = true;
5679
});
5780
const envs: NativeEnvInfo[] = [];
5881
let discovered = createDeferred();
59-
const disposable = result.discovered((data) => envs.push(data));
60-
82+
const disposable = result.discovered((data) => {
83+
envs.push(data);
84+
discovered.resolve();
85+
});
6186
do {
62-
await Promise.race([result.completed, discovered.promise]);
87+
if (!envs.length) {
88+
await Promise.race([result.completed, discovered.promise]);
89+
}
6390
if (envs.length) {
6491
const dataToSend = [...envs];
6592
envs.length = 0;
@@ -69,27 +96,22 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
6996
}
7097
if (!completed) {
7198
discovered = createDeferred();
72-
envs.length = 0;
7399
}
74100
} while (!completed);
75-
76101
disposable.dispose();
77102
}
78103

79104
// eslint-disable-next-line class-methods-use-this
80-
private start(): { completed: Promise<void>; discovered: Event<NativeEnvInfo> } {
81-
const discovered = new EventEmitter<NativeEnvInfo>();
82-
const completed = createDeferred<void>();
105+
private start(): rpc.MessageConnection {
83106
const proc = ch.spawn(NATIVE_LOCATOR, ['server'], { env: process.env });
84107
const disposables: Disposable[] = [];
85108
// jsonrpc package cannot handle messages coming through too quicly.
86109
// Lets handle the messages and close the stream only when
87110
// we have got the exit event.
88111
const readable = new PassThrough();
89112
proc.stdout.pipe(readable, { end: false });
90-
let err = '';
91113
proc.stderr.on('data', (data) => {
92-
err += data.toString();
114+
const err = data.toString();
93115
traceError('Native Python Finder', err);
94116
});
95117
const writable = new PassThrough();
@@ -105,17 +127,10 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
105127
disposables.push(
106128
connection,
107129
disposeStreams,
108-
discovered,
109130
connection.onError((ex) => {
110131
disposeStreams.dispose();
111132
traceError('Error in Native Python Finder', ex);
112133
}),
113-
connection.onNotification('environment', (data: NativeEnvInfo) => {
114-
discovered.fire(data);
115-
}),
116-
// connection.onNotification((method: string, data: any) => {
117-
// console.log(method, data);
118-
// }),
119134
connection.onNotification('log', (data: NativeLog) => {
120135
switch (data.level) {
121136
case 'info':
@@ -135,7 +150,6 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
135150
}
136151
}),
137152
connection.onClose(() => {
138-
completed.resolve();
139153
disposables.forEach((d) => d.dispose());
140154
}),
141155
{
@@ -152,19 +166,86 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
152166
);
153167

154168
connection.listen();
155-
connection
156-
.sendRequest<number>('refresh', {
157-
// Send configuration information to the Python finder.
158-
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
159-
conda_executable: undefined,
160-
})
161-
.then((durationInMilliSeconds: number) => {
162-
completed.resolve();
163-
traceInfo(`Native Python Finder took ${durationInMilliSeconds}ms to complete.`);
164-
})
165-
.catch((ex) => traceError('Error in Native Python Finder', ex));
166-
167-
return { completed: completed.promise, discovered: discovered.event };
169+
this._register(Disposable.from(...disposables));
170+
return connection;
171+
}
172+
173+
private doRefresh(): { completed: Promise<void>; discovered: Event<NativeEnvInfo> } {
174+
const disposable = this._register(new DisposableStore());
175+
const discovered = disposable.add(new EventEmitter<NativeEnvInfo>());
176+
const completed = createDeferred<void>();
177+
const pendingPromises: Promise<void>[] = [];
178+
179+
const notifyUponCompletion = () => {
180+
const initialCount = pendingPromises.length;
181+
Promise.all(pendingPromises)
182+
.then(() => {
183+
if (initialCount === pendingPromises.length) {
184+
completed.resolve();
185+
} else {
186+
setTimeout(notifyUponCompletion, 0);
187+
}
188+
})
189+
.catch(noop);
190+
};
191+
const trackPromiseAndNotifyOnCompletion = (promise: Promise<void>) => {
192+
pendingPromises.push(promise);
193+
notifyUponCompletion();
194+
};
195+
196+
disposable.add(
197+
this.connection.onNotification('environment', (data: NativeEnvInfo) => {
198+
// We know that in the Python extension if either Version of Prefix is not provided by locator
199+
// Then we end up resolving the information.
200+
// Lets do that here,
201+
// This is a hack, as the other part of the code that resolves the version information
202+
// doesn't work as expected, as its still a WIP.
203+
if (data.executable && (!data.version || !data.prefix)) {
204+
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
205+
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
206+
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
207+
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
208+
const promise = this.connection
209+
.sendRequest<{ duration: number; environment: NativeEnvInfo }>('resolve', {
210+
executable: data.executable,
211+
})
212+
.then(({ environment, duration }) => {
213+
traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
214+
discovered.fire(environment);
215+
})
216+
.catch((ex) => traceError(`Error in Resolving Python Environment ${data}`, ex));
217+
trackPromiseAndNotifyOnCompletion(promise);
218+
} else {
219+
discovered.fire(data);
220+
}
221+
}),
222+
);
223+
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+
);
232+
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+
})
241+
.then(({ duration }) => traceInfo(`Native Python Finder completed in ${duration}ms`))
242+
.catch((ex) => traceError('Error in Native Python Finder', ex)),
243+
);
244+
completed.promise.finally(() => disposable.dispose());
245+
return {
246+
completed: completed.promise,
247+
discovered: discovered.event,
248+
};
168249
}
169250
}
170251

src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,11 @@ export async function resolveBasicEnv(env: BasicEnvInfo): Promise<PythonEnvInfo>
5959
const resolvedEnv = await resolverForKind(env);
6060
resolvedEnv.searchLocation = getSearchLocation(resolvedEnv, searchLocation);
6161
resolvedEnv.source = uniq(resolvedEnv.source.concat(source ?? []));
62-
if (getOSType() === OSType.Windows && resolvedEnv.source?.includes(PythonEnvSource.WindowsRegistry)) {
62+
if (
63+
!env.identifiedUsingNativeLocator &&
64+
getOSType() === OSType.Windows &&
65+
resolvedEnv.source?.includes(PythonEnvSource.WindowsRegistry)
66+
) {
6367
// We can update env further using information we can get from the Windows registry.
6468
await updateEnvUsingRegistry(resolvedEnv);
6569
}
@@ -75,9 +79,11 @@ export async function resolveBasicEnv(env: BasicEnvInfo): Promise<PythonEnvInfo>
7579
resolvedEnv.executable.ctime = ctime;
7680
resolvedEnv.executable.mtime = mtime;
7781
}
78-
const type = await getEnvType(resolvedEnv);
79-
if (type) {
80-
resolvedEnv.type = type;
82+
if (!env.identifiedUsingNativeLocator) {
83+
const type = await getEnvType(resolvedEnv);
84+
if (type) {
85+
resolvedEnv.type = type;
86+
}
8187
}
8288
return resolvedEnv;
8389
}
@@ -147,7 +153,7 @@ async function resolveGloballyInstalledEnv(env: BasicEnvInfo): Promise<PythonEnv
147153
const { executablePath } = env;
148154
let version;
149155
try {
150-
version = parseVersionFromExecutable(executablePath);
156+
version = env.identifiedUsingNativeLocator ? env.version : parseVersionFromExecutable(executablePath);
151157
} catch {
152158
version = UNKNOWN_PYTHON_VERSION;
153159
}
@@ -169,7 +175,7 @@ async function resolveSimpleEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
169175
const { executablePath, kind } = env;
170176
const envInfo = buildEnvInfo({
171177
kind,
172-
version: await getPythonVersionFromPath(executablePath),
178+
version: env.identifiedUsingNativeLocator ? env.version : await getPythonVersionFromPath(executablePath),
173179
executable: executablePath,
174180
sysPrefix: env.envPath,
175181
location: env.envPath,

src/client/pythonEnvironments/index.ts

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -197,15 +197,20 @@ function watchRoots(args: WatchRootsArgs): IDisposable {
197197
}
198198

199199
function createWorkspaceLocator(ext: ExtensionState): WorkspaceLocators {
200-
const locators = new WorkspaceLocators(watchRoots, [
201-
(root: vscode.Uri) => [
202-
new WorkspaceVirtualEnvironmentLocator(root.fsPath),
203-
new PoetryLocator(root.fsPath),
204-
new HatchLocator(root.fsPath),
205-
new CustomWorkspaceLocator(root.fsPath),
206-
],
207-
// Add an ILocator factory func here for each kind of workspace-rooted locator.
208-
]);
200+
const locators = new WorkspaceLocators(
201+
watchRoots,
202+
useNativeLocator()
203+
? []
204+
: [
205+
(root: vscode.Uri) => [
206+
new WorkspaceVirtualEnvironmentLocator(root.fsPath),
207+
new PoetryLocator(root.fsPath),
208+
new HatchLocator(root.fsPath),
209+
new CustomWorkspaceLocator(root.fsPath),
210+
],
211+
// Add an ILocator factory func here for each kind of workspace-rooted locator.
212+
],
213+
);
209214
ext.disposables.push(locators);
210215
return locators;
211216
}

0 commit comments

Comments
 (0)