Skip to content

Commit

Permalink
Adopt native locator resolve methods (microsoft#23631)
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne authored Jun 17, 2024
1 parent 0835c7e commit 8a87e1d
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 48 deletions.
147 changes: 114 additions & 33 deletions src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@ 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 { DisposableBase } from '../../../../common/utils/resourceLifecycle';
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';

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

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

Expand All @@ -48,18 +52,41 @@ interface NativeLog {
}

class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGlobalPythonFinder {
private readonly connection: rpc.MessageConnection;

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

public async resolve(executable: string): Promise<NativeEnvInfo> {
const { environment, duration } = await this.connection.sendRequest<{
duration: number;
environment: NativeEnvInfo;
}>('resolve', {
executable,
});

traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
return environment;
}

async *refresh(_paths: Uri[]): AsyncIterable<NativeEnvInfo> {
const result = this.start();
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));

const disposable = result.discovered((data) => {
envs.push(data);
discovered.resolve();
});
do {
await Promise.race([result.completed, discovered.promise]);
if (!envs.length) {
await Promise.race([result.completed, discovered.promise]);
}
if (envs.length) {
const dataToSend = [...envs];
envs.length = 0;
Expand All @@ -69,27 +96,22 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
}
if (!completed) {
discovered = createDeferred();
envs.length = 0;
}
} while (!completed);

disposable.dispose();
}

// eslint-disable-next-line class-methods-use-this
private start(): { completed: Promise<void>; discovered: Event<NativeEnvInfo> } {
const discovered = new EventEmitter<NativeEnvInfo>();
const completed = createDeferred<void>();
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.
// Lets handle the messages and close the stream only when
// we have got the exit event.
const readable = new PassThrough();
proc.stdout.pipe(readable, { end: false });
let err = '';
proc.stderr.on('data', (data) => {
err += data.toString();
const err = data.toString();
traceError('Native Python Finder', err);
});
const writable = new PassThrough();
Expand All @@ -105,17 +127,10 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
disposables.push(
connection,
disposeStreams,
discovered,
connection.onError((ex) => {
disposeStreams.dispose();
traceError('Error in Native Python Finder', ex);
}),
connection.onNotification('environment', (data: NativeEnvInfo) => {
discovered.fire(data);
}),
// connection.onNotification((method: string, data: any) => {
// console.log(method, data);
// }),
connection.onNotification('log', (data: NativeLog) => {
switch (data.level) {
case 'info':
Expand All @@ -135,7 +150,6 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
}
}),
connection.onClose(() => {
completed.resolve();
disposables.forEach((d) => d.dispose());
}),
{
Expand All @@ -152,19 +166,86 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
);

connection.listen();
connection
.sendRequest<number>('refresh', {
// Send configuration information to the Python finder.
search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
conda_executable: undefined,
})
.then((durationInMilliSeconds: number) => {
completed.resolve();
traceInfo(`Native Python Finder took ${durationInMilliSeconds}ms to complete.`);
})
.catch((ex) => traceError('Error in Native Python Finder', ex));

return { completed: completed.promise, discovered: discovered.event };
this._register(Disposable.from(...disposables));
return connection;
}

private doRefresh(): { completed: Promise<void>; discovered: Event<NativeEnvInfo> } {
const disposable = this._register(new DisposableStore());
const discovered = disposable.add(new EventEmitter<NativeEnvInfo>());
const completed = createDeferred<void>();
const pendingPromises: Promise<void>[] = [];

const notifyUponCompletion = () => {
const initialCount = pendingPromises.length;
Promise.all(pendingPromises)
.then(() => {
if (initialCount === pendingPromises.length) {
completed.resolve();
} else {
setTimeout(notifyUponCompletion, 0);
}
})
.catch(noop);
};
const trackPromiseAndNotifyOnCompletion = (promise: Promise<void>) => {
pendingPromises.push(promise);
notifyUponCompletion();
};

disposable.add(
this.connection.onNotification('environment', (data: NativeEnvInfo) => {
// We know that in the Python extension if either Version of Prefix is not provided by locator
// Then we end up resolving the information.
// Lets do that here,
// This is a hack, as the other part of the code that resolves the version information
// doesn't work as expected, as its still a WIP.
if (data.executable && (!data.version || !data.prefix)) {
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
// HACK = TEMPORARY WORK AROUND, TO GET STUFF WORKING
const promise = this.connection
.sendRequest<{ duration: number; environment: NativeEnvInfo }>('resolve', {
executable: data.executable,
})
.then(({ environment, duration }) => {
traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
discovered.fire(environment);
})
.catch((ex) => traceError(`Error in Resolving Python Environment ${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,
})
.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,
};
}
}

Expand Down
18 changes: 12 additions & 6 deletions src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,11 @@ export async function resolveBasicEnv(env: BasicEnvInfo): Promise<PythonEnvInfo>
const resolvedEnv = await resolverForKind(env);
resolvedEnv.searchLocation = getSearchLocation(resolvedEnv, searchLocation);
resolvedEnv.source = uniq(resolvedEnv.source.concat(source ?? []));
if (getOSType() === OSType.Windows && resolvedEnv.source?.includes(PythonEnvSource.WindowsRegistry)) {
if (
!env.identifiedUsingNativeLocator &&
getOSType() === OSType.Windows &&
resolvedEnv.source?.includes(PythonEnvSource.WindowsRegistry)
) {
// We can update env further using information we can get from the Windows registry.
await updateEnvUsingRegistry(resolvedEnv);
}
Expand All @@ -75,9 +79,11 @@ export async function resolveBasicEnv(env: BasicEnvInfo): Promise<PythonEnvInfo>
resolvedEnv.executable.ctime = ctime;
resolvedEnv.executable.mtime = mtime;
}
const type = await getEnvType(resolvedEnv);
if (type) {
resolvedEnv.type = type;
if (!env.identifiedUsingNativeLocator) {
const type = await getEnvType(resolvedEnv);
if (type) {
resolvedEnv.type = type;
}
}
return resolvedEnv;
}
Expand Down Expand Up @@ -147,7 +153,7 @@ async function resolveGloballyInstalledEnv(env: BasicEnvInfo): Promise<PythonEnv
const { executablePath } = env;
let version;
try {
version = parseVersionFromExecutable(executablePath);
version = env.identifiedUsingNativeLocator ? env.version : parseVersionFromExecutable(executablePath);
} catch {
version = UNKNOWN_PYTHON_VERSION;
}
Expand All @@ -169,7 +175,7 @@ async function resolveSimpleEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
const { executablePath, kind } = env;
const envInfo = buildEnvInfo({
kind,
version: await getPythonVersionFromPath(executablePath),
version: env.identifiedUsingNativeLocator ? env.version : await getPythonVersionFromPath(executablePath),
executable: executablePath,
sysPrefix: env.envPath,
location: env.envPath,
Expand Down
23 changes: 14 additions & 9 deletions src/client/pythonEnvironments/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,15 +197,20 @@ function watchRoots(args: WatchRootsArgs): IDisposable {
}

function createWorkspaceLocator(ext: ExtensionState): WorkspaceLocators {
const locators = new WorkspaceLocators(watchRoots, [
(root: vscode.Uri) => [
new WorkspaceVirtualEnvironmentLocator(root.fsPath),
new PoetryLocator(root.fsPath),
new HatchLocator(root.fsPath),
new CustomWorkspaceLocator(root.fsPath),
],
// Add an ILocator factory func here for each kind of workspace-rooted locator.
]);
const locators = new WorkspaceLocators(
watchRoots,
useNativeLocator()
? []
: [
(root: vscode.Uri) => [
new WorkspaceVirtualEnvironmentLocator(root.fsPath),
new PoetryLocator(root.fsPath),
new HatchLocator(root.fsPath),
new CustomWorkspaceLocator(root.fsPath),
],
// Add an ILocator factory func here for each kind of workspace-rooted locator.
],
);
ext.disposables.push(locators);
return locators;
}
Expand Down

0 comments on commit 8a87e1d

Please sign in to comment.