Updates to use new Python locator Api (microsoft#23832)

Author: DonJayamanne

src/client/common/utils/async.ts

@@ -232,6 +232,17 @@ export async function flattenIterator<T>(iterator: IAsyncIterator<T>): Promise<T
     return results;
 }
 
+/**
+ * Get everything yielded by the iterable.
+ */
+export async function flattenIterable<T>(iterableItem: AsyncIterable<T>): Promise<T[]> {
+    const results: T[] = [];
+    for await (const item of iterableItem) {
+        results.push(item);
+    }
+    return results;
+}
+
 /**
  * Wait for a condition to be fulfilled within a timeout.
  *

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

@@ -18,6 +18,7 @@ import { getUserHomeDir } from '../../../../common/utils/platform';
 import { createLogOutputChannel } from '../../../../common/vscodeApis/windowApis';
 import { PythonEnvKind } from '../../info';
 import { sendNativeTelemetry, NativePythonTelemetry } from './nativePythonTelemetry';
+import { traceError } from '../../../../logging';
 
 const untildify = require('untildify');
 
@@ -29,7 +30,7 @@ export interface NativeEnvInfo {
     displayName?: string;
     name?: string;
     executable?: string;
-    kind?: string;
+    kind?: PythonEnvironmentKind;
     version?: string;
     prefix?: string;
     manager?: NativeEnvManagerInfo;
@@ -41,12 +42,38 @@ export interface NativeEnvInfo {
     symlinks?: string[];
 }
 
+export enum PythonEnvironmentKind {
+    Conda = 'Conda',
+    Homebrew = 'Homebrew',
+    Pyenv = 'Pyenv',
+    GlobalPaths = 'GlobalPaths',
+    PyenvVirtualEnv = 'PyenvVirtualEnv',
+    Pipenv = 'Pipenv',
+    Poetry = 'Poetry',
+    MacPythonOrg = 'MacPythonOrg',
+    MacCommandLineTools = 'MacCommandLineTools',
+    LinuxGlobal = 'LinuxGlobal',
+    MacXCode = 'MacXCode',
+    Venv = 'Venv',
+    VirtualEnv = 'VirtualEnv',
+    VirtualEnvWrapper = 'VirtualEnvWrapper',
+    WindowsStore = 'WindowsStore',
+    WindowsRegistry = 'WindowsRegistry',
+}
+
 export interface NativeEnvManagerInfo {
     tool: string;
     executable: string;
     version?: string;
 }
 
+export function isNativeInfoEnvironment(info: NativeEnvInfo | NativeEnvManagerInfo): info is NativeEnvInfo {
+    if ((info as NativeEnvManagerInfo).tool) {
+        return false;
+    }
+    return true;
+}
+
 export type NativeCondaInfo = {
     canSpawnConda: boolean;
     userProvidedEnvFound?: boolean;
@@ -58,12 +85,62 @@ export type NativeCondaInfo = {
 };
 
 export interface NativePythonFinder extends Disposable {
+    /**
+     * Refresh the list of python environments.
+     * Returns an async iterable that can be used to iterate over the list of python environments.
+     * Internally this will take all of the current workspace folders and search for python environments.
+     *
+     * If a Uri is provided, then it will search for python environments in that location (ignoring workspaces).
+     * Uri can be a file or a folder.
+     * If a PythonEnvironmentKind is provided, then it will search for python environments of that kind (ignoring workspaces).
+     */
+    refresh(options?: PythonEnvironmentKind | Uri[]): AsyncIterable<NativeEnvInfo | NativeEnvManagerInfo>;
+    /**
+     * Will spawn the provided Python executable and return information about the environment.
+     * @param executable
+     */
     resolve(executable: string): Promise<NativeEnvInfo>;
-    refresh(): AsyncIterable<NativeEnvInfo>;
-    categoryToKind(category?: string): PythonEnvKind;
-    logger(): LogOutputChannel;
+    categoryToKind(category?: PythonEnvironmentKind): PythonEnvKind;
+    /**
+     * Used only for telemetry.
+     */
     getCondaInfo(): Promise<NativeCondaInfo>;
-    find(searchPath: string): Promise<NativeEnvInfo[]>;
+}
+
+const mapping = new Map<PythonEnvironmentKind, PythonEnvKind>([
+    [PythonEnvironmentKind.Conda, PythonEnvKind.Conda],
+    [PythonEnvironmentKind.GlobalPaths, PythonEnvKind.OtherGlobal],
+    [PythonEnvironmentKind.Pyenv, PythonEnvKind.Pyenv],
+    [PythonEnvironmentKind.PyenvVirtualEnv, PythonEnvKind.Pyenv],
+    [PythonEnvironmentKind.Pipenv, PythonEnvKind.Pipenv],
+    [PythonEnvironmentKind.Poetry, PythonEnvKind.Poetry],
+    [PythonEnvironmentKind.VirtualEnv, PythonEnvKind.VirtualEnv],
+    [PythonEnvironmentKind.VirtualEnvWrapper, PythonEnvKind.VirtualEnvWrapper],
+    [PythonEnvironmentKind.Venv, PythonEnvKind.Venv],
+    [PythonEnvironmentKind.WindowsRegistry, PythonEnvKind.System],
+    [PythonEnvironmentKind.WindowsStore, PythonEnvKind.MicrosoftStore],
+    [PythonEnvironmentKind.Homebrew, PythonEnvKind.System],
+    [PythonEnvironmentKind.LinuxGlobal, PythonEnvKind.System],
+    [PythonEnvironmentKind.MacCommandLineTools, PythonEnvKind.System],
+    [PythonEnvironmentKind.MacPythonOrg, PythonEnvKind.System],
+    [PythonEnvironmentKind.MacXCode, PythonEnvKind.System],
+]);
+
+export function categoryToKind(category?: PythonEnvironmentKind, logger?: LogOutputChannel): PythonEnvKind {
+    if (!category) {
+        return PythonEnvKind.Unknown;
+    }
+    const kind = mapping.get(category);
+    if (kind) {
+        return kind;
+    }
+
+    if (logger) {
+        logger.error(`Unknown Python Environment category '${category}' from Native Locator.`);
+    } else {
+        traceError(`Unknown Python Environment category '${category}' from Native Locator.`);
+    }
+    return PythonEnvKind.Unknown;
 }
 
 interface NativeLog {
@@ -94,47 +171,11 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
         return environment;
     }
 
-    categoryToKind(category?: string): PythonEnvKind {
-        if (!category) {
-            return PythonEnvKind.Unknown;
-        }
-        switch (category.toLowerCase()) {
-            case 'conda':
-                return PythonEnvKind.Conda;
-            case 'system':
-            case 'homebrew':
-            case 'macpythonorg':
-            case 'maccommandlinetools':
-            case 'macxcode':
-            case 'windowsregistry':
-            case 'linuxglobal':
-                return PythonEnvKind.System;
-            case 'globalpaths':
-                return PythonEnvKind.OtherGlobal;
-            case 'pyenv':
-                return PythonEnvKind.Pyenv;
-            case 'poetry':
-                return PythonEnvKind.Poetry;
-            case 'pipenv':
-                return PythonEnvKind.Pipenv;
-            case 'pyenvvirtualenv':
-                return PythonEnvKind.VirtualEnv;
-            case 'venv':
-                return PythonEnvKind.Venv;
-            case 'virtualenv':
-                return PythonEnvKind.VirtualEnv;
-            case 'virtualenvwrapper':
-                return PythonEnvKind.VirtualEnvWrapper;
-            case 'windowsstore':
-                return PythonEnvKind.MicrosoftStore;
-            default: {
-                this.outputChannel.info(`Unknown Python Environment category '${category}' from Native Locator.`);
-                return PythonEnvKind.Unknown;
-            }
-        }
+    categoryToKind(category?: PythonEnvironmentKind): PythonEnvKind {
+        return categoryToKind(category, this.outputChannel);
     }
 
-    async *refresh(): AsyncIterable<NativeEnvInfo> {
+    async *refresh(options?: PythonEnvironmentKind | Uri[]): AsyncIterable<NativeEnvInfo> {
         if (this.firstRefreshResults) {
             // If this is the first time we are refreshing,
             // Then get the results from the first refresh.
@@ -143,12 +184,12 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
             this.firstRefreshResults = undefined;
             yield* results;
         } else {
-            const result = this.doRefresh();
+            const result = this.doRefresh(options);
             let completed = false;
             void result.completed.finally(() => {
                 completed = true;
             });
-            const envs: NativeEnvInfo[] = [];
+            const envs: (NativeEnvInfo | NativeEnvManagerInfo)[] = [];
             let discovered = createDeferred();
             const disposable = result.discovered((data) => {
                 envs.push(data);
@@ -173,10 +214,6 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
         }
     }
 
-    logger(): LogOutputChannel {
-        return this.outputChannel;
-    }
-
     refreshFirstTime() {
         const result = this.doRefresh();
         const completed = createDeferredFrom(result.completed);
@@ -283,9 +320,11 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
         return connection;
     }
 
-    private doRefresh(): { completed: Promise<void>; discovered: Event<NativeEnvInfo> } {
+    private doRefresh(
+        options?: PythonEnvironmentKind | Uri[],
+    ): { completed: Promise<void>; discovered: Event<NativeEnvInfo | NativeEnvManagerInfo> } {
         const disposable = this._register(new DisposableStore());
-        const discovered = disposable.add(new EventEmitter<NativeEnvInfo>());
+        const discovered = disposable.add(new EventEmitter<NativeEnvInfo | NativeEnvManagerInfo>());
         const completed = createDeferred<void>();
         const pendingPromises: Promise<void>[] = [];
 
@@ -306,6 +345,8 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
             notifyUponCompletion();
         };
 
+        // Assumption is server will ensure there's only one refresh at a time.
+        // Perhaps we should have a request Id or the like to map the results back to the `refresh` request.
         disposable.add(
             this.connection.onNotification('environment', (data: NativeEnvInfo) => {
                 this.outputChannel.info(`Discovered env: ${data.executable || data.prefix}`);
@@ -334,11 +375,28 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativePytho
                 }
             }),
         );
+        disposable.add(
+            this.connection.onNotification('manager', (data: NativeEnvManagerInfo) => {
+                this.outputChannel.info(`Discovered manager: (${data.tool}) ${data.executable}`);
+                discovered.fire(data);
+            }),
+        );
 
+        type RefreshOptions = {
+            searchKind?: PythonEnvironmentKind;
+            searchPaths?: string[];
+        };
+
+        const refreshOptions: RefreshOptions = {};
+        if (options && Array.isArray(options) && options.length > 0) {
+            refreshOptions.searchPaths = options.map((item) => item.fsPath);
+        } else if (options && typeof options === 'string') {
+            refreshOptions.searchKind = options;
+        }
         trackPromiseAndNotifyOnCompletion(
             this.configure().then(() =>
                 this.connection
-                    .sendRequest<{ duration: number }>('refresh')
+                    .sendRequest<{ duration: number }>('refresh', refreshOptions)
                     .then(({ duration }) => this.outputChannel.info(`Refresh completed in ${duration}ms`))
                     .catch((ex) => this.outputChannel.error('Refresh error', ex)),
             ),

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

@@ -2,9 +2,9 @@
 // Licensed under the MIT License.
 
 import * as fsPath from 'path';
-import { Event, EventEmitter, workspace } from 'vscode';
+import { Event, EventEmitter, Uri, workspace } from 'vscode';
 import '../../../../common/extensions';
-import { createDeferred, Deferred } from '../../../../common/utils/async';
+import { createDeferred, Deferred, flattenIterable } from '../../../../common/utils/async';
 import { StopWatch } from '../../../../common/utils/stopWatch';
 import { traceError, traceInfo, traceVerbose } from '../../../../logging';
 import { sendTelemetryEvent } from '../../../../telemetry';
@@ -25,7 +25,12 @@ import {
 import { getQueryFilter } from '../../locatorUtils';
 import { PythonEnvCollectionChangedEvent, PythonEnvsWatcher } from '../../watcher';
 import { IEnvsCollectionCache } from './envsCollectionCache';
-import { getNativePythonFinder, NativeEnvInfo, NativePythonFinder } from '../common/nativePythonFinder';
+import {
+    getNativePythonFinder,
+    isNativeInfoEnvironment,
+    NativeEnvInfo,
+    NativePythonFinder,
+} from '../common/nativePythonFinder';
 import { pathExists } from '../../../../common/platform/fs-paths';
 import { noop } from '../../../../common/utils/misc';
 import { parseVersion } from '../../info/pythonVersion';
@@ -294,16 +299,18 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
         const executablesFoundByNativeLocator = new Set<string>();
         const nativeStopWatch = new StopWatch();
         for await (const data of this.nativeFinder.refresh()) {
-            nativeEnvs.push(data);
-            if (data.executable) {
-                // Lowercase for purposes of comparison (safe).
-                executablesFoundByNativeLocator.add(data.executable.toLowerCase());
-            } else if (data.prefix) {
+            if (isNativeInfoEnvironment(data)) {
+                nativeEnvs.push(data);
+                if (data.executable) {
+                    // Lowercase for purposes of comparison (safe).
+                    executablesFoundByNativeLocator.add(data.executable.toLowerCase());
+                } else if (data.prefix) {
+                    // Lowercase for purposes of comparison (safe).
+                    executablesFoundByNativeLocator.add(data.prefix.toLowerCase());
+                }
                 // Lowercase for purposes of comparison (safe).
-                executablesFoundByNativeLocator.add(data.prefix.toLowerCase());
+                (data.symlinks || []).forEach((exe) => executablesFoundByNativeLocator.add(exe.toLowerCase()));
             }
-            // Lowercase for purposes of comparison (safe).
-            (data.symlinks || []).forEach((exe) => executablesFoundByNativeLocator.add(exe.toLowerCase()));
         }
         const nativeDuration = nativeStopWatch.elapsedTime;
         void this.sendNativeLocatorTelemetry(nativeEnvs);
@@ -980,11 +987,11 @@ async function getCondaTelemetry(
 
         if (condaTelemetry.condaRootPrefixFoundInInfoNotInNative) {
             // Verify we are able to discover this environment as a conda env using native finder.
-            const rootPrefixEnvs = await nativeFinder.find(rootPrefix);
+            const rootPrefixEnvs = await flattenIterable(nativeFinder.refresh([Uri.file(rootPrefix)]));
             // Did we find an env with the same prefix?
-            const rootPrefixEnv = rootPrefixEnvs.find(
-                (e) => fsPath.normalize(e.prefix || '').toLowerCase() === rootPrefix.toLowerCase(),
-            );
+            const rootPrefixEnv = rootPrefixEnvs
+                .filter(isNativeInfoEnvironment)
+                .find((e) => fsPath.normalize(e.prefix || '').toLowerCase() === rootPrefix.toLowerCase());
             condaTelemetry.condaRootPrefixEnvsAfterFind = rootPrefixEnvs.length;
             condaTelemetry.condaRootPrefixFoundInInfoAfterFind = !!rootPrefixEnv;
             condaTelemetry.condaRootPrefixFoundInInfoAfterFindKind = rootPrefixEnv?.kind;
@@ -1019,11 +1026,11 @@ async function getCondaTelemetry(
 
         if (condaTelemetry.condaDefaultPrefixFoundInInfoNotInNative) {
             // Verify we are able to discover this environment as a conda env using native finder.
-            const defaultPrefixEnvs = await nativeFinder.find(defaultPrefix);
+            const defaultPrefixEnvs = await flattenIterable(nativeFinder.refresh([Uri.file(defaultPrefix)]));
             // Did we find an env with the same prefix?
-            const defaultPrefixEnv = defaultPrefixEnvs.find(
-                (e) => fsPath.normalize(e.prefix || '').toLowerCase() === defaultPrefix.toLowerCase(),
-            );
+            const defaultPrefixEnv = defaultPrefixEnvs
+                .filter(isNativeInfoEnvironment)
+                .find((e) => fsPath.normalize(e.prefix || '').toLowerCase() === defaultPrefix.toLowerCase());
             condaTelemetry.condaDefaultPrefixEnvsAfterFind = defaultPrefixEnvs.length;
             condaTelemetry.condaDefaultPrefixFoundInInfoAfterFind = !!defaultPrefixEnv;
             condaTelemetry.condaDefaultPrefixFoundInInfoAfterFindKind = defaultPrefixEnv?.kind;