Add WorkspaceLocators.

Author: ericsnowcurrently

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

@@ -3,7 +3,7 @@ import { Disposable, Event, EventEmitter, Uri } from 'vscode';
 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,
@@ -18,13 +18,112 @@ import {
     WORKSPACE_VIRTUAL_ENV_SERVICE
 } from '../../../interpreter/contracts';
 import { IServiceContainer } from '../../../ioc/types';
+import { PythonEnvInfo } from '../../base/info';
+import { FullLocator, ILocator, 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');
 
+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 FullLocator {
+    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<PythonEnvInfo>(iterators);
+    }
+
+    private addRoot(root: Uri) {
+        // Drop the old one, 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;
+        // Hook up the watchers.
+        locator.onChanged((e) => 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();
+    }
+}
+
+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.
  */