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

ericsnowcurrently · web-flow · commit b6ee6c61d11a · 2020-09-09T13:22:06.000-06:00
diff --git a/src/client/common/utils/async.ts b/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.
  */
diff --git a/src/client/pythonEnvironments/discovery/locators/index.ts b/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.
  */
diff --git a/src/test/pythonEnvironments/base/common.ts b/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) {
diff --git a/src/test/pythonEnvironments/discovery/locators/index.unit.test.ts b/src/test/pythonEnvironments/discovery/locators/index.unit.test.ts
