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

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;