Add an extension-specific locator that will wrap the low-level locators. (microsoft#13780)

Author: Eric Snow

src/client/common/utils/async.ts

@@ -137,7 +137,7 @@ async function getNext<T>(it: AsyncIterator<T, T | void>, indexMaybe?: number):
 }
 
 // tslint:disable-next-line:promise-must-complete no-empty
-const NEVER: Promise<unknown> = new Promise(() => {});
+export const NEVER: Promise<unknown> = new Promise(() => {});
 
 /**
  * Yield everything produced by the given iterators as soon as each is ready.
@@ -181,6 +181,43 @@ export async function* chain<T>(
  }
 }
 
+/**
+ * Map the async function onto the items and yield the results.
+ *
+ * @param items - the items to map onto and iterate
+ * @param func - the async function to apply for each item
+ * @param race - if `true` (the default) then results are yielded
+ * potentially out of order, as soon as each is ready
+ */
+export async function* mapToIterator<T, R = T>(
+ items: T[],
+ func: (item: T) => Promise<R>,
+ race = true
+): AsyncIterator<R, void> {
+ if (race) {
+ const iterators = items.map((item) => {
+ async function* generator() {
+ yield func(item);
+ }
+ return generator();
+ });
+ yield* iterable(chain(iterators));
+ } else {
+ yield* items.map(func);
+ }
+}
+
+/**
+ * Convert an iterator into an iterable, if it isn't one already.
+ */
+export function iterable<T>(iterator: AsyncIterator<T, void>): AsyncIterableIterator<T> {
+ const it = iterator as AsyncIterableIterator<T>;
+ if (it[Symbol.asyncIterator] === undefined) {
+ it[Symbol.asyncIterator] = () => it;
+ }
+ return it;
+}
+
 /**
  * Get everything yielded by the iterator.
  */

src/client/pythonEnvironments/discovery/locators/index.ts

@@ -5,7 +5,7 @@ import {
 import { traceDecorators } from '../../../common/logger';
 import { IPlatformService } from '../../../common/platform/types';
 import { IDisposableRegistry } from '../../../common/types';
-import { createDeferred, Deferred } from '../../../common/utils/async';
+import { chain, createDeferred, Deferred } from '../../../common/utils/async';
 import { OSType } from '../../../common/utils/platform';
 import {
  CONDA_ENV_FILE_SERVICE,
@@ -20,13 +20,157 @@ import {
  WORKSPACE_VIRTUAL_ENV_SERVICE,
 } from '../../../interpreter/contracts';
 import { IServiceContainer } from '../../../ioc/types';
+import { PythonEnvInfo } from '../../base/info';
+import { ILocator, Locator, NOOP_ITERATOR, PythonEnvsIterator, PythonLocatorQuery } from '../../base/locator';
+import { DisableableLocator, Locators } from '../../base/locators';
 import { PythonEnvironment } from '../../info';
 import { isHiddenInterpreter } from './services/interpreterFilter';
 import { GetInterpreterLocatorOptions } from './types';
 
 // tslint:disable-next-line:no-require-imports no-var-requires
 const flatten = require('lodash/flatten') as typeof import('lodash/flatten');
 
+/**
+ * A wrapper around all locators used by the extension.
+ */
+export class ExtensionLocators extends Locators {
+ constructor(
+ // These are expected to be low-level locators (e.g. system).
+ nonWorkspace: ILocator[],
+ // This is expected to be a locator wrapping any found in
+ // the workspace (i.e. WorkspaceLocators).
+ workspace: ILocator
+ ) {
+ super([...nonWorkspace, workspace]);
+ }
+}
+
+type WorkspaceLocatorFactory = (root: Uri) => ILocator[];
+
+interface IWorkspaceFolders {
+ readonly roots: ReadonlyArray<Uri>;
+ readonly onAdded: Event<Uri>;
+ readonly onRemoved: Event<Uri>;
+}
+
+type RootURI = string;
+
+/**
+ * The collection of all workspace-specific locators used by the extension.
+ *
+ * The factories are used to produce the locators for each workspace folder.
+ */
+export class WorkspaceLocators extends Locator {
+ private readonly locators: Record<RootURI, DisableableLocator> = {};
+ private readonly roots: Record<RootURI, Uri> = {};
+ constructor(
+ // used to produce the per-root locators:
+ private readonly factories: WorkspaceLocatorFactory[]
+ ) {
+ super();
+ }
+
+ /**
+ * Activate the locator.
+ *
+ * @param folders - the info used to keep track of the workspace folders
+ */
+ public activate(folders: IWorkspaceFolders) {
+ for (const root of folders.roots) {
+ this.addRoot(root);
+ }
+ folders.onAdded((root: Uri) => this.addRoot(root));
+ folders.onRemoved((root: Uri) => this.removeRoot(root));
+ }
+
+ public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+ const iterators = Object.keys(this.locators).map((key) => {
+ if (query?.searchLocations) {
+ const root = this.roots[key];
+ if (!matchURI(root, ...query.searchLocations)) {
+ return NOOP_ITERATOR;
+ }
+ }
+ // The query matches or was not location-specific.
+ const locator = this.locators[key];
+ return locator.iterEnvs(query);
+ });
+ return chain(iterators);
+ }
+
+ public async resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
+ if (typeof env !== 'string' && env.searchLocation) {
+ const rootLocator = this.locators[env.searchLocation.toString()];
+ if (rootLocator) {
+ return rootLocator.resolveEnv(env);
+ }
+ }
+ // Fall back to checking all the roots.
+ for (const key of Object.keys(this.locators)) {
+ const resolved = await this.locators[key].resolveEnv(env);
+ if (resolved !== undefined) {
+ return resolved;
+ }
+ }
+ return undefined;
+ }
+
+ private addRoot(root: Uri) {
+ // Drop the old one, if necessary.
+ this.removeRoot(root);
+ // Create the root's locator, wrapping each factory-generated locator.
+ const locators: ILocator[] = [];
+ for (const create of this.factories) {
+ locators.push(...create(root));
+ }
+ const locator = new DisableableLocator(new Locators(locators));
+ // Cache it.
+ const key = root.toString();
+ this.locators[key] = locator;
+ this.roots[key] = root;
+ this.emitter.fire({ searchLocation: root });
+ // Hook up the watchers.
+ locator.onChanged((e) => {
+ if (e.searchLocation === undefined) {
+ e.searchLocation = root;
+ }
+ this.emitter.fire(e);
+ });
+ }
+
+ private removeRoot(root: Uri) {
+ const key = root.toString();
+ const locator = this.locators[key];
+ if (locator === undefined) {
+ return;
+ }
+ delete this.locators[key];
+ delete this.roots[key];
+ locator.disable();
+ this.emitter.fire({ searchLocation: root });
+ }
+}
+
+/**
+ * Determine if the given URI matches one of the candidates.
+ *
+ * The scheme must match, as well as path. The path must match exactly
+ * or the URI must be a parent of one of the candidates.
+ */
+function matchURI(uri: Uri, ...candidates: Uri[]): boolean {
+ const uriPath = uri.path.endsWith('/') ? uri.path : `{uri.path}/`;
+ for (const candidate of candidates) {
+ if (candidate.scheme === uri.scheme) {
+ if (candidate.path === uri.path) {
+ return true;
+ } else if (candidate.path.startsWith(uriPath)) {
+ return true;
+ }
+ }
+ }
+ return false;
+}
+
 /**
  * Facilitates locating Python interpreters.
  */

src/test/pythonEnvironments/base/common.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { createDeferred, flattenIterator } from '../../../client/common/utils/async';
+import { createDeferred, flattenIterator, iterable, mapToIterator } from '../../../client/common/utils/async';
 import { Architecture } from '../../../client/common/utils/platform';
 import { EMPTY_VERSION, parseBasicVersionInfo } from '../../../client/common/utils/version';
 import {
@@ -122,14 +122,24 @@ export class SimpleLocator extends Locator {
  if (callbacks?.before !== undefined) {
  await callbacks.before;
  }
- //yield* envs;
- for (const env of envs) {
- if (callbacks?.beforeEach !== undefined) {
- await callbacks.beforeEach(env);
+ if (callbacks?.beforeEach !== undefined) {
+ // The results will likely come in a different order.
+ const mapped = mapToIterator(envs, async (env) => {
+ await callbacks.beforeEach!(env);
+ return env;
+ });
+ for await (const env of iterable(mapped)) {
+ yield env;
+ if (callbacks?.afterEach !== undefined) {
+ await callbacks.afterEach(env);
+ }
  }
- yield env;
- if (callbacks?.afterEach !== undefined) {
- await callbacks.afterEach(env);
+ } else {
+ for (const env of envs) {
+ yield env;
+ if (callbacks?.afterEach !== undefined) {
+ await callbacks.afterEach(env);
+ }
  }
  }
  if (callbacks?.after!== undefined) {