Add API to contribute a new locator

Author: Kartik Raj

src/client/pythonEnvironments/base/locators.ts

@@ -3,8 +3,8 @@
 
 import { chain } from '../../common/utils/async';
 import { Disposables } from '../../common/utils/resourceLifecycle';
-import { PythonEnvInfo } from './info';
 import {
+    BasicEnvInfo,
     ILocator,
     IPythonEnvsIterator,
     isProgressEvent,
@@ -59,10 +59,10 @@ export function combineIterators<I>(iterators: IPythonEnvsIterator<I>[]): IPytho
  *
  * Events and iterator results are combined.
  */
-export class Locators<I = PythonEnvInfo> extends PythonEnvsWatchers implements ILocator<I> {
+export class Locators<I = BasicEnvInfo> extends PythonEnvsWatchers implements ILocator<I> {
     constructor(
         // The locators will be watched as well as iterated.
-        private readonly locators: ReadonlyArray<ILocator<I>>,
+        private locators: ILocator<I>[],
     ) {
         super(locators);
     }
@@ -71,4 +71,9 @@ export class Locators<I = PythonEnvInfo> extends PythonEnvsWatchers implements I
         const iterators = this.locators.map((loc) => loc.iterEnvs(query));
         return combineIterators(iterators);
     }
+
+    public addLocator(locator: ILocator<I>): void {
+        this.locators = [...this.locators, locator];
+        this.addWatcher(locator);
+    }
 }

src/client/pythonEnvironments/base/locators/wrappers.ts

@@ -3,37 +3,48 @@
 
 // eslint-disable-next-line max-classes-per-file
 import { Uri } from 'vscode';
+import { ILocatorClass } from '../../../apiTypes';
 import { IDisposable } from '../../../common/types';
 import { iterEmpty } from '../../../common/utils/async';
 import { getURIFilter } from '../../../common/utils/misc';
 import { Disposables } from '../../../common/utils/resourceLifecycle';
+import { CustomLocator } from '../../converter';
 import { PythonEnvInfo } from '../info';
-import { ILocator, IPythonEnvsIterator, PythonLocatorQuery } from '../locator';
+import { BasicEnvInfo, ILocator, IPythonEnvsIterator, PythonLocatorQuery } from '../locator';
 import { combineIterators, Locators } from '../locators';
 import { LazyResourceBasedLocator } from './common/resourceBasedLocator';
 
 /**
  * A wrapper around all locators used by the extension.
  */
 
-export class ExtensionLocators<I = PythonEnvInfo> extends Locators<I> {
+export class ExtensionLocators extends Locators<BasicEnvInfo> {
     constructor(
         // These are expected to be low-level locators (e.g. system).
-        private readonly nonWorkspace: ILocator<I>[],
+        private nonWorkspace: ILocator<BasicEnvInfo>[],
         // This is expected to be a locator wrapping any found in
         // the workspace (i.e. WorkspaceLocators).
-        private readonly workspace: ILocator<I>,
+        private workspace: WorkspaceLocators,
     ) {
         super([...nonWorkspace, workspace]);
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I> {
-        const iterators: IPythonEnvsIterator<I>[] = [this.workspace.iterEnvs(query)];
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
+        const iterators: IPythonEnvsIterator<BasicEnvInfo>[] = [this.workspace.iterEnvs(query)];
         if (!query?.searchLocations?.doNotIncludeNonRooted) {
             iterators.push(...this.nonWorkspace.map((loc) => loc.iterEnvs(query)));
         }
         return combineIterators(iterators);
     }
+
+    public addNewLocator(LocatorClass: ILocatorClass, isWorkspace: boolean): void {
+        if (isWorkspace) {
+            this.workspace.addNewLocator(LocatorClass);
+        }
+        if (!isWorkspace) {
+            this.nonWorkspace = [...this.nonWorkspace, new CustomLocator(new LocatorClass())];
+        }
+    }
 }
 type WorkspaceLocatorFactoryResult<I> = ILocator<I> & Partial<IDisposable>;
 type WorkspaceLocatorFactory<I = PythonEnvInfo> = (root: Uri) => WorkspaceLocatorFactoryResult<I>[];
@@ -52,12 +63,15 @@ type WatchRootsFunc = (args: WatchRootsArgs) => IDisposable;
  * The factories are used to produce the locators for each workspace folder.
  */
 
-export class WorkspaceLocators<I = PythonEnvInfo> extends LazyResourceBasedLocator<I> {
-    private readonly locators: Record<RootURI, [ILocator<I>, IDisposable]> = {};
+export class WorkspaceLocators extends LazyResourceBasedLocator<BasicEnvInfo> {
+    private readonly locators: Record<RootURI, [Locators<BasicEnvInfo>, IDisposable]> = {};
 
     private readonly roots: Record<RootURI, Uri> = {};
 
-    constructor(private readonly watchRoots: WatchRootsFunc, private readonly factories: WorkspaceLocatorFactory<I>[]) {
+    constructor(
+        private readonly watchRoots: WatchRootsFunc,
+        private readonly factory: WorkspaceLocatorFactory<BasicEnvInfo>,
+    ) {
         super();
     }
 
@@ -69,7 +83,7 @@ export class WorkspaceLocators<I = PythonEnvInfo> extends LazyResourceBasedLocat
         roots.forEach((root) => this.removeRoot(root));
     }
 
-    protected doIterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I> {
+    protected doIterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
         const iterators = Object.keys(this.locators).map((key) => {
             if (query?.searchLocations !== undefined) {
                 const root = this.roots[key];
@@ -78,7 +92,7 @@ export class WorkspaceLocators<I = PythonEnvInfo> extends LazyResourceBasedLocat
                 // Ignore any requests for global envs.
                 if (!query.searchLocations.roots.some(filter)) {
                     // This workspace folder did not match the query, so skip it!
-                    return iterEmpty<I>();
+                    return iterEmpty<BasicEnvInfo>();
                 }
             }
             // The query matches or was not location-specific.
@@ -107,15 +121,13 @@ export class WorkspaceLocators<I = PythonEnvInfo> extends LazyResourceBasedLocat
 
     private addRoot(root: Uri): void {
         // Create the root's locator, wrapping each factory-generated locator.
-        const locators: ILocator<I>[] = [];
+        const locators: ILocator<BasicEnvInfo>[] = [];
         const disposables = new Disposables();
-        this.factories.forEach((create) => {
-            create(root).forEach((loc) => {
-                locators.push(loc);
-                if (loc.dispose !== undefined) {
-                    disposables.push(loc as IDisposable);
-                }
-            });
+        this.factory(root).forEach((loc) => {
+            locators.push(loc);
+            if (loc.dispose !== undefined) {
+                disposables.push(loc as IDisposable);
+            }
         });
         const locator = new Locators(locators);
         // Cache it.
@@ -133,6 +145,16 @@ export class WorkspaceLocators<I = PythonEnvInfo> extends LazyResourceBasedLocat
         );
     }
 
+    public addNewLocator(LocatorClass: ILocatorClass): void {
+        Object.keys(this.roots).forEach((key) => {
+            const root = this.roots[key];
+            const newLocator = new LocatorClass(root.fsPath);
+            const convertedLocator: ILocator<BasicEnvInfo> = new CustomLocator(newLocator);
+            const [locators] = this.locators[key];
+            locators.addLocator(convertedLocator);
+        });
+    }
+
     private removeRoot(root: Uri): void {
         const key = root.toString();
         const found = this.locators[key];

src/client/pythonEnvironments/base/watchers.ts

@@ -30,4 +30,10 @@ export class PythonEnvsWatchers implements IPythonEnvsWatcher, IDisposable {
     public async dispose(): Promise<void> {
         await this.disposables.dispose();
     }
+
+    // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
+    protected addWatcher(w: IPythonEnvsWatcher) {
+        const disposable = w.onChanged((e) => this.watcher.fire(e));
+        this.disposables.push(disposable);
+    }
 }

src/client/pythonEnvironments/index.ts

@@ -106,7 +106,7 @@ async function createLocator(
     // This is shared.
 ): Promise<IDiscoveryAPI> {
     // Create the low-level locators.
-    let locators: ILocator<BasicEnvInfo> = new ExtensionLocators<BasicEnvInfo>(
+    let locators: ILocator<BasicEnvInfo> = new ExtensionLocators(
         // Here we pull the locators together.
         createNonWorkspaceLocators(ext),
         createWorkspaceLocator(ext),
@@ -177,10 +177,10 @@ function watchRoots(args: WatchRootsArgs): IDisposable {
     });
 }
 
-function createWorkspaceLocator(ext: ExtensionState): WorkspaceLocators<BasicEnvInfo> {
-    const locators = new WorkspaceLocators<BasicEnvInfo>(watchRoots, [
-        (root: vscode.Uri) => [new WorkspaceVirtualEnvironmentLocator(root.fsPath), new PoetryLocator(root.fsPath)],
-        // Add an ILocator factory func here for each kind of workspace-rooted locator.
+function createWorkspaceLocator(ext: ExtensionState): WorkspaceLocators {
+    const locators = new WorkspaceLocators(watchRoots, (root: vscode.Uri) => [
+        new WorkspaceVirtualEnvironmentLocator(root.fsPath),
+        new PoetryLocator(root.fsPath),
     ]);
     ext.disposables.push(locators);
     return locators;