Make lookup for new environments detected faster (#19922)

Closes #19900

Trigger a scoped refresh instead of a full refresh when an environment
gets created, preferably only looking for the environment which gets
created.

Author: Kartik Raj

src/client/pythonEnvironments/base/locator.ts

@@ -7,7 +7,6 @@ import { Event, Uri } from 'vscode';
 import { IAsyncIterableIterator, iterEmpty } from '../../common/utils/async';
 import { PythonEnvInfo, PythonEnvKind, PythonEnvSource } from './info';
 import {
-    BasicPythonEnvsChangedEvent,
     IPythonEnvsWatcher,
     PythonEnvCollectionChangedEvent,
     PythonEnvsChangedEvent,
@@ -128,6 +127,14 @@ export type PythonLocatorQuery = BasicPythonLocatorQuery & {
      * If provided, results should be limited to within these locations.
      */
     searchLocations?: SearchLocations;
+    /**
+     * If provided, results should be limited envs provided by these locators.
+     */
+    providerId?: string;
+    /**
+     * If provided, results area limited to this env.
+     */
+    envPath?: string;
 };
 
 type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;
@@ -153,8 +160,8 @@ export type BasicEnvInfo = {
  * events emitted via `onChanged` do not need to provide information
  * for the specific environments that changed.
  */
-export interface ILocator<I = PythonEnvInfo, E extends BasicPythonEnvsChangedEvent = PythonEnvsChangedEvent>
-    extends IPythonEnvsWatcher<E> {
+export interface ILocator<I = PythonEnvInfo, E = PythonEnvsChangedEvent> extends IPythonEnvsWatcher<E> {
+    readonly providerId: string;
     /**
      * Iterate over the enviroments known tos this locator.
      *
@@ -174,6 +181,8 @@ export interface ILocator<I = PythonEnvInfo, E extends BasicPythonEnvsChangedEve
     iterEnvs(query?: QueryForEvent<E>): IPythonEnvsIterator<I>;
 }
 
+export type ICompositeLocator<I = PythonEnvInfo, E = PythonEnvsChangedEvent> = Omit<ILocator<I, E>, 'providerId'>;
+
 interface IResolver {
     /**
      * Find as much info about the given Python environment as possible.
@@ -184,7 +193,7 @@ interface IResolver {
     resolveEnv(path: string): Promise<PythonEnvInfo | undefined>;
 }
 
-export interface IResolvingLocator<I = PythonEnvInfo> extends IResolver, ILocator<I> {}
+export interface IResolvingLocator<I = PythonEnvInfo> extends IResolver, ICompositeLocator<I> {}
 
 export interface GetRefreshEnvironmentsOptions {
     /**
@@ -234,7 +243,7 @@ export interface IDiscoveryAPI {
     resolveEnv(path: string): Promise<PythonEnvInfo | undefined>;
 }
 
-interface IEmitter<E extends PythonEnvsChangedEvent> {
+interface IEmitter<E> {
     fire(e: E): void;
 }
 
@@ -250,10 +259,11 @@ interface IEmitter<E extends PythonEnvsChangedEvent> {
  * should be used.  Only in low-level cases should you consider using
  * `BasicPythonEnvsChangedEvent`.
  */
-abstract class LocatorBase<I = PythonEnvInfo, E extends BasicPythonEnvsChangedEvent = PythonEnvsChangedEvent>
-    implements ILocator<I, E> {
+abstract class LocatorBase<I = PythonEnvInfo, E = PythonEnvsChangedEvent> implements ILocator<I, E> {
     public readonly onChanged: Event<E>;
 
+    public abstract readonly providerId: string;
+
     protected readonly emitter: IEmitter<E>;
 
     constructor(watcher: IPythonEnvsWatcher<E> & IEmitter<E>) {

src/client/pythonEnvironments/base/locators.ts

@@ -5,6 +5,7 @@ import { chain } from '../../common/utils/async';
 import { Disposables } from '../../common/utils/resourceLifecycle';
 import { PythonEnvInfo } from './info';
 import {
+    ICompositeLocator,
     ILocator,
     IPythonEnvsIterator,
     isProgressEvent,
@@ -59,12 +60,15 @@ 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 = PythonEnvInfo> extends PythonEnvsWatchers implements ICompositeLocator<I> {
+    public readonly providerId: string;
+
     constructor(
         // The locators will be watched as well as iterated.
         private readonly locators: ReadonlyArray<ILocator<I>>,
     ) {
         super(locators);
+        this.providerId = locators.map((loc) => loc.providerId).join('+');
     }
 
     public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I> {

src/client/pythonEnvironments/base/locators/common/resourceBasedLocator.ts

@@ -5,8 +5,9 @@ import { IDisposable } from '../../../../common/types';
 import { createDeferred, Deferred } from '../../../../common/utils/async';
 import { Disposables } from '../../../../common/utils/resourceLifecycle';
 import { traceError } from '../../../../logging';
-import { PythonEnvInfo } from '../../info';
-import { IPythonEnvsIterator, Locator, PythonLocatorQuery } from '../../locator';
+import { arePathsSame } from '../../../common/externalDependencies';
+import { getEnvPath } from '../../info/env';
+import { BasicEnvInfo, IPythonEnvsIterator, Locator, PythonLocatorQuery } from '../../locator';
 
 /**
  * A base locator class that manages the lifecycle of resources.
@@ -20,7 +21,7 @@ import { IPythonEnvsIterator, Locator, PythonLocatorQuery } from '../../locator'
  *
  * Otherwise it will leak (and we have no leak detection).
  */
-export abstract class LazyResourceBasedLocator<I = PythonEnvInfo> extends Locator<I> implements IDisposable {
+export abstract class LazyResourceBasedLocator extends Locator<BasicEnvInfo> implements IDisposable {
     protected readonly disposables = new Disposables();
 
     // This will be set only once we have to create necessary resources
@@ -42,15 +43,29 @@ export abstract class LazyResourceBasedLocator<I = PythonEnvInfo> extends Locato
         await this.disposables.dispose();
     }
 
-    public async *iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I> {
+    public async *iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
         await this.activate();
-        yield* this.doIterEnvs(query);
+        const iterator = this.doIterEnvs(query);
+        if (query?.envPath) {
+            let result = await iterator.next();
+            while (!result.done) {
+                const currEnv = result.value;
+                const { path } = getEnvPath(currEnv.executablePath, currEnv.envPath);
+                if (arePathsSame(path, query.envPath)) {
+                    yield currEnv;
+                    break;
+                }
+                result = await iterator.next();
+            }
+        } else {
+            yield* iterator;
+        }
     }
 
     /**
      * The subclass implementation of iterEnvs().
      */
-    protected abstract doIterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I>;
+    protected abstract doIterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo>;
 
     /**
      * This is where subclasses get their resources ready.

src/client/pythonEnvironments/base/locators/composite/envsCollectionService.ts

@@ -57,7 +57,9 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
     constructor(private readonly cache: IEnvsCollectionCache, private readonly locator: IResolvingLocator) {
         super();
         this.locator.onChanged((event) => {
-            const query = undefined; // We can also form a query based on the event, but skip that for simplicity.
+            const query: PythonLocatorQuery | undefined = event.providerId
+                ? { providerId: event.providerId, envPath: event.envPath }
+                : undefined; // We can also form a query based on the event, but skip that for simplicity.
             let scheduledRefresh = this.scheduledRefreshesPerQuery.get(query);
             // If there is no refresh scheduled for the query, start a new one.
             if (!scheduledRefresh) {

src/client/pythonEnvironments/base/locators/composite/envsReducer.ts

@@ -9,6 +9,7 @@ import { areSameEnv } from '../../info/env';
 import { getPrioritizedEnvKinds } from '../../info/envKind';
 import {
     BasicEnvInfo,
+    ICompositeLocator,
     ILocator,
     IPythonEnvsIterator,
     isProgressEvent,
@@ -22,7 +23,7 @@ import { PythonEnvsChangedEvent } from '../../watcher';
 /**
  * Combines duplicate environments received from the incoming locator into one and passes on unique environments
  */
-export class PythonEnvsReducer implements ILocator<BasicEnvInfo> {
+export class PythonEnvsReducer implements ICompositeLocator<BasicEnvInfo> {
     public get onChanged(): Event<PythonEnvsChangedEvent> {
         return this.parentLocator.onChanged;
     }

src/client/pythonEnvironments/base/locators/composite/envsResolver.ts

@@ -10,7 +10,7 @@ import { getEnvPath, setEnvDisplayString } from '../../info/env';
 import { InterpreterInformation } from '../../info/interpreter';
 import {
     BasicEnvInfo,
-    ILocator,
+    ICompositeLocator,
     IPythonEnvsIterator,
     IResolvingLocator,
     isProgressEvent,
@@ -35,7 +35,7 @@ export class PythonEnvsResolver implements IResolvingLocator {
     }
 
     constructor(
-        private readonly parentLocator: ILocator<BasicEnvInfo>,
+        private readonly parentLocator: ICompositeLocator<BasicEnvInfo>,
         private readonly environmentInfoService: IEnvironmentInfoService,
     ) {}
 

src/client/pythonEnvironments/base/locators/lowLevel/condaLocator.ts

@@ -7,7 +7,9 @@ import { Conda, getCondaEnvironmentsTxt } from '../../../common/environmentManag
 import { traceError, traceVerbose } from '../../../../logging';
 import { FSWatchingLocator } from './fsWatchingLocator';
 
-export class CondaEnvironmentLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class CondaEnvironmentLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'conda-envs';
+
     public constructor() {
         super(
             () => getCondaEnvironmentsTxt(),

src/client/pythonEnvironments/base/locators/lowLevel/customVirtualEnvLocator.ts

@@ -78,7 +78,9 @@ async function getVirtualEnvKind(interpreterPath: string): Promise<PythonEnvKind
 /**
  * Finds and resolves custom virtual environments that users have provided.
  */
-export class CustomVirtualEnvironmentLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class CustomVirtualEnvironmentLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'custom-virtual-envs';
+
     constructor() {
         super(getCustomVirtualEnvDirs, getVirtualEnvKind, {
             // Note detecting kind of virtual env depends on the file structure around the

src/client/pythonEnvironments/base/locators/lowLevel/filesLocator.ts

@@ -14,7 +14,9 @@ type GetExecutablesFunc = () => AsyncIterableIterator<string>;
 /**
  * A naive locator the wraps a function that finds Python executables.
  */
-class FoundFilesLocator implements ILocator<BasicEnvInfo> {
+abstract class FoundFilesLocator implements ILocator<BasicEnvInfo> {
+    public abstract readonly providerId: string;
+
     public readonly onChanged: Event<PythonEnvsChangedEvent>;
 
     protected readonly watcher = new PythonEnvsWatcher();
@@ -45,6 +47,8 @@ type GetDirExecutablesFunc = (dir: string) => AsyncIterableIterator<string>;
  * A locator for executables in a single directory.
  */
 export class DirFilesLocator extends FoundFilesLocator {
+    public readonly providerId: string;
+
     constructor(
         dirname: string,
         defaultKind: PythonEnvKind,
@@ -53,6 +57,7 @@ export class DirFilesLocator extends FoundFilesLocator {
         source?: PythonEnvSource[],
     ) {
         super(defaultKind, () => getExecutables(dirname), source);
+        this.providerId = `dir-files-${dirname}`;
     }
 }
 

src/client/pythonEnvironments/base/locators/lowLevel/fsWatchingLocator.ts

@@ -13,7 +13,7 @@ import {
     resolvePythonExeGlobs,
     watchLocationForPythonBinaries,
 } from '../../../common/pythonBinariesWatcher';
-import { PythonEnvInfo, PythonEnvKind } from '../../info';
+import { PythonEnvKind } from '../../info';
 import { LazyResourceBasedLocator } from '../common/resourceBasedLocator';
 
 export enum FSWatcherKind {
@@ -80,7 +80,7 @@ type FileWatchOptions = {
  *
  * Subclasses can call `this.emitter.fire()` * to emit events.
  */
-export abstract class FSWatchingLocator<I = PythonEnvInfo> extends LazyResourceBasedLocator<I> {
+export abstract class FSWatchingLocator extends LazyResourceBasedLocator {
     constructor(
         /**
          * Location(s) to watch for python binaries.
@@ -135,7 +135,7 @@ export abstract class FSWatchingLocator<I = PythonEnvInfo> extends LazyResourceB
             this.disposables.push(
                 watchLocationForPattern(path.dirname(root), path.basename(root), () => {
                     traceVerbose('Detected change in file: ', root, 'initiating a refresh');
-                    this.emitter.fire({});
+                    this.emitter.fire({ providerId: this.providerId });
                 }),
             );
             return;
@@ -161,7 +161,7 @@ export abstract class FSWatchingLocator<I = PythonEnvInfo> extends LazyResourceB
             //        |__ python  <--- executable
             const searchLocation = Uri.file(opts.searchLocation ?? path.dirname(getEnvironmentDirFromPath(executable)));
             traceVerbose('Fired event ', JSON.stringify({ type, kind, searchLocation }), 'from locator');
-            this.emitter.fire({ type, kind, searchLocation });
+            this.emitter.fire({ type, kind, searchLocation, providerId: this.providerId, envPath: executable });
         };
 
         const globs = resolvePythonExeGlobs(

src/client/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvronmentLocator.ts

@@ -82,7 +82,9 @@ async function getVirtualEnvKind(interpreterPath: string): Promise<PythonEnvKind
 /**
  * Finds and resolves virtual environments created in known global locations.
  */
-export class GlobalVirtualEnvironmentLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class GlobalVirtualEnvironmentLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'global-virtual-env';
+
     constructor(private readonly searchDepth?: number) {
         super(getGlobalVirtualEnvDirs, getVirtualEnvKind, {
             // Note detecting kind of virtual env depends on the file structure around the

src/client/pythonEnvironments/base/locators/lowLevel/microsoftStoreLocator.ts

@@ -65,7 +65,9 @@ export async function getMicrosoftStorePythonExes(): Promise<string[]> {
     return [];
 }
 
-export class MicrosoftStoreLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class MicrosoftStoreLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'microsoft-store';
+
     private readonly kind: PythonEnvKind = PythonEnvKind.MicrosoftStore;
 
     constructor() {

src/client/pythonEnvironments/base/locators/lowLevel/poetryLocator.ts

@@ -60,7 +60,9 @@ async function getVirtualEnvKind(interpreterPath: string): Promise<PythonEnvKind
 /**
  * Finds and resolves virtual environments created using poetry.
  */
-export class PoetryLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class PoetryLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'poetry';
+
     public constructor(private readonly root: string) {
         super(
             () => getRootVirtualEnvDir(root),

src/client/pythonEnvironments/base/locators/lowLevel/posixKnownPathsLocator.ts

@@ -12,6 +12,8 @@ import { isMacDefaultPythonPath } from './macDefaultLocator';
 import { traceError } from '../../../../logging';
 
 export class PosixKnownPathsLocator extends Locator<BasicEnvInfo> {
+    public readonly providerId = 'posixKnownPaths';
+
     private kind: PythonEnvKind = PythonEnvKind.OtherGlobal;
 
     public iterEnvs(): IPythonEnvsIterator<BasicEnvInfo> {

src/client/pythonEnvironments/base/locators/lowLevel/pyenvLocator.ts

@@ -35,7 +35,9 @@ async function* getPyenvEnvironments(): AsyncIterableIterator<BasicEnvInfo> {
     }
 }
 
-export class PyenvLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class PyenvLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'pyenv';
+
     constructor() {
         super(getPyenvVersionsDir, async () => PythonEnvKind.Pyenv);
     }

src/client/pythonEnvironments/base/locators/lowLevel/windowsKnownPathsLocator.ts

@@ -24,6 +24,8 @@ import { DirFilesLocator } from './filesLocator';
  * it for changes.
  */
 export class WindowsPathEnvVarLocator implements ILocator<BasicEnvInfo>, IDisposable {
+    public readonly providerId: string = 'windows-path-env-var-locator';
+
     public readonly onChanged: Event<PythonEnvsChangedEvent>;
 
     private readonly locators: Locators<BasicEnvInfo>;
@@ -93,6 +95,7 @@ function getDirFilesLocator(
         yield* await getEnvs(locator.iterEnvs(query));
     }
     return {
+        providerId: locator.providerId,
         iterEnvs,
         dispose,
         onChanged: locator.onChanged,

src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts

@@ -7,6 +7,8 @@ import { getRegistryInterpreters } from '../../../common/windowsUtils';
 import { traceError } from '../../../../logging';
 
 export class WindowsRegistryLocator extends Locator<BasicEnvInfo> {
+    public readonly providerId: string = 'windows-registry';
+
     // eslint-disable-next-line class-methods-use-this
     public iterEnvs(): IPythonEnvsIterator<BasicEnvInfo> {
         const iterator = async function* () {

src/client/pythonEnvironments/base/locators/lowLevel/workspaceVirtualEnvLocator.ts

@@ -50,7 +50,9 @@ async function getVirtualEnvKind(interpreterPath: string): Promise<PythonEnvKind
 /**
  * Finds and resolves virtual environments created in workspace roots.
  */
-export class WorkspaceVirtualEnvironmentLocator extends FSWatchingLocator<BasicEnvInfo> {
+export class WorkspaceVirtualEnvironmentLocator extends FSWatchingLocator {
+    public readonly providerId: string = 'workspaceVirtualEnvLocator';
+
     public constructor(private readonly root: string) {
         super(
             () => getWorkspaceVirtualEnvDirs(this.root),