Add an API to report progress of environment discovery in two phases (#19125)

* Add progress API

* Remove old API in favor of this

* Revert "Remove old API in favor of this"

This reverts commit 9543501.

* Revert "Revert "Remove old API in favor of this""

This reverts commit 5bbea78.

* Fix

* Remove old API impl

* Add to proposed API

* Add get refresh promise options

* Change getter to function

* Translate progress events to progress promises

* Fix combining iterators

* Fix tests

* Add test for resolver

* Add test for reducer

* Fixes

* Add tests for environment collection service

* News entry

* Add another test if a query is provided

* Add comments and clarify

* Fix bug

* Add clarifying comment

Author: Kartik Raj

news/1 Enhancements/19103.md

@@ -0,0 +1 @@
+Added a proposed API to report progress of environment discovery in two phases.

src/client/apiTypes.ts

@@ -5,6 +5,7 @@ import { Event, Uri } from 'vscode';
 import { Resource } from './common/types';
 import { IDataViewerDataProvider, IJupyterUriProvider } from './jupyter/types';
 import { EnvPathType, PythonEnvKind } from './pythonEnvironments/base/info';
+import { GetRefreshEnvironmentsOptions, ProgressNotificationEvent } from './pythonEnvironments/base/locator';
 
 /*
  * Do not introduce any breaking changes to this API.
@@ -203,11 +204,17 @@ export interface IProposedExtensionAPI {
          *                    is triggered.
          */
         refreshEnvironment(options?: RefreshEnvironmentsOptions): Promise<EnvPathType[] | undefined>;
+        /**
+         * Tracks discovery progress for current list of known environments, i.e when it starts, finishes or any other relevant
+         * stage. Note the progress for a particular query is currently not tracked or reported, this only indicates progress of
+         * the entire collection.
+         */
+        readonly onRefreshProgress: Event<ProgressNotificationEvent>;
         /**
          * Returns a promise for the ongoing refresh. Returns `undefined` if there are no active
          * refreshes going on.
          */
-        getRefreshPromise(): Promise<void> | undefined;
+        getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined;
         /**
          * This event is triggered when the known environment list changes, like when a environment
          * is found, existing environment is removed, or some details changed on an environment.

src/client/interpreter/contracts.ts

@@ -3,7 +3,7 @@ import { CodeLensProvider, ConfigurationTarget, Disposable, Event, TextDocument,
 import { FileChangeType } from '../common/platform/fileSystemWatcher';
 import { Resource } from '../common/types';
 import { PythonEnvSource } from '../pythonEnvironments/base/info';
-import { PythonLocatorQuery } from '../pythonEnvironments/base/locator';
+import { ProgressNotificationEvent, PythonLocatorQuery } from '../pythonEnvironments/base/locator';
 import { CondaEnvironmentInfo } from '../pythonEnvironments/common/environmentManagers/conda';
 import { EnvironmentType, PythonEnvironment } from '../pythonEnvironments/info';
 
@@ -16,9 +16,9 @@ export type PythonEnvironmentsChangedEvent = {
 
 export const IComponentAdapter = Symbol('IComponentAdapter');
 export interface IComponentAdapter {
-    readonly onRefreshStart: Event<void>;
+    readonly onProgress: Event<ProgressNotificationEvent>;
     triggerRefresh(query?: PythonLocatorQuery & { clearCache?: boolean }, trigger?: 'auto' | 'ui'): Promise<void>;
-    readonly refreshPromise: Promise<void> | undefined;
+    getRefreshPromise(): Promise<void> | undefined;
     readonly onChanged: Event<PythonEnvironmentsChangedEvent>;
     // VirtualEnvPrompt
     onDidCreate(resource: Resource, callback: () => void): Disposable;
@@ -63,7 +63,6 @@ export interface ICondaService {
 
 export const IInterpreterService = Symbol('IInterpreterService');
 export interface IInterpreterService {
-    readonly onRefreshStart: Event<void>;
     triggerRefresh(query?: PythonLocatorQuery & { clearCache?: boolean }, trigger?: 'auto' | 'ui'): Promise<void>;
     readonly refreshPromise: Promise<void> | undefined;
     readonly onDidChangeInterpreters: Event<PythonEnvironmentsChangedEvent>;

src/client/interpreter/display/progressDisplay.ts

@@ -12,6 +12,7 @@ import { IDisposableRegistry } from '../../common/types';
 import { createDeferred, Deferred } from '../../common/utils/async';
 import { Interpreters } from '../../common/utils/localize';
 import { traceDecoratorVerbose } from '../../logging';
+import { ProgressReportStage } from '../../pythonEnvironments/base/locator';
 import { IComponentAdapter } from '../contracts';
 
 // The parts of IComponentAdapter used here.
@@ -30,11 +31,14 @@ export class InterpreterLocatorProgressStatubarHandler implements IExtensionSing
     ) {}
 
     public async activate(): Promise<void> {
-        this.pyenvs.onRefreshStart(
-            () => {
-                this.showProgress();
-                if (this.pyenvs.refreshPromise) {
-                    this.pyenvs.refreshPromise.then(() => this.hideProgress());
+        this.pyenvs.onProgress(
+            (event) => {
+                if (event.stage === ProgressReportStage.discoveryStarted) {
+                    this.showProgress();
+                    const refreshPromise = this.pyenvs.getRefreshPromise();
+                    if (refreshPromise) {
+                        refreshPromise.then(() => this.hideProgress());
+                    }
                 }
             },
             this,

src/client/interpreter/interpreterService.ts

@@ -40,10 +40,6 @@ export class InterpreterService implements Disposable, IInterpreterService {
         return this.pyenvs.hasInterpreters(filter);
     }
 
-    public get onRefreshStart(): Event<void> {
-        return this.pyenvs.onRefreshStart;
-    }
-
     public triggerRefresh(
         query?: PythonLocatorQuery & { clearCache?: boolean },
         trigger?: 'auto' | 'ui',
@@ -52,7 +48,7 @@ export class InterpreterService implements Disposable, IInterpreterService {
     }
 
     public get refreshPromise(): Promise<void> | undefined {
-        return this.pyenvs.refreshPromise;
+        return this.pyenvs.getRefreshPromise();
     }
 
     public get onDidChangeInterpreter(): Event<void> {

src/client/proposedApi.ts

@@ -16,7 +16,7 @@ import { IInterpreterService } from './interpreter/contracts';
 import { IServiceContainer } from './ioc/types';
 import { PythonEnvInfo } from './pythonEnvironments/base/info';
 import { getEnvPath } from './pythonEnvironments/base/info/env';
-import { IDiscoveryAPI } from './pythonEnvironments/base/locator';
+import { GetRefreshEnvironmentsOptions, IDiscoveryAPI } from './pythonEnvironments/base/locator';
 
 const onDidInterpretersChangedEvent = new EventEmitter<EnvironmentsChangedParams[]>();
 export function reportInterpretersChanged(e: EnvironmentsChangedParams[]): void {
@@ -106,12 +106,13 @@ export function buildProposedApi(
                 const paths = discoveryApi.getEnvs().map((e) => getEnvPath(e.executable.filename, e.location));
                 return Promise.resolve(paths);
             },
-            getRefreshPromise(): Promise<void> | undefined {
-                return discoveryApi.refreshPromise;
+            getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined {
+                return discoveryApi.getRefreshPromise(options);
             },
             onDidChangeExecutionDetails: interpreterService.onDidChangeInterpreterConfiguration,
             onDidEnvironmentsChanged: onDidInterpretersChangedEvent.event,
             onDidActiveEnvironmentChanged: onDidActiveInterpreterChangedEvent.event,
+            onRefreshProgress: discoveryApi.onProgress,
         },
     };
     return proposed;

src/client/pythonEnvironments/api.ts

@@ -1,11 +1,17 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import { Event } from 'vscode';
 import { StopWatch } from '../common/utils/stopWatch';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { getEnvPath } from './base/info/env';
-import { IDiscoveryAPI, PythonLocatorQuery } from './base/locator';
+import {
+    GetRefreshEnvironmentsOptions,
+    IDiscoveryAPI,
+    ProgressNotificationEvent,
+    PythonLocatorQuery,
+} from './base/locator';
 
 export type GetLocatorFunc = () => Promise<IDiscoveryAPI>;
 
@@ -26,12 +32,12 @@ class PythonEnvironments implements IDiscoveryAPI {
         this.locator = await this.getLocator();
     }
 
-    public get onRefreshStart() {
-        return this.locator.onRefreshStart;
+    public get onProgress(): Event<ProgressNotificationEvent> {
+        return this.locator.onProgress;
     }
 
-    public get refreshPromise() {
-        return this.locator.refreshPromise;
+    public getRefreshPromise(options?: GetRefreshEnvironmentsOptions) {
+        return this.locator.getRefreshPromise(options);
     }
 
     public get onChanged() {

src/client/pythonEnvironments/base/locator.ts

@@ -64,7 +64,23 @@ export interface IPythonEnvsIterator<I = PythonEnvInfo> extends IAsyncIterableIt
      * If this property is not provided then it means the iterator does
      * not support updates.
      */
-    onUpdated?: Event<PythonEnvUpdatedEvent<I> | null>;
+    onUpdated?: Event<PythonEnvUpdatedEvent<I> | ProgressNotificationEvent>;
+}
+
+export enum ProgressReportStage {
+    discoveryStarted = 'discoveryStarted',
+    allPathsDiscovered = 'allPathsDiscovered',
+    discoveryFinished = 'discoveryFinished',
+}
+
+export type ProgressNotificationEvent = {
+    stage: ProgressReportStage;
+};
+
+export function isProgressEvent<I = PythonEnvInfo>(
+    event: PythonEnvUpdatedEvent<I> | ProgressNotificationEvent,
+): event is ProgressNotificationEvent {
+    return 'stage' in event;
 }
 
 /**
@@ -170,11 +186,20 @@ interface IResolver {
 
 export interface IResolvingLocator<I = PythonEnvInfo> extends IResolver, ILocator<I> {}
 
+export interface GetRefreshEnvironmentsOptions {
+    /**
+     * Get refresh promise which resolves once the following stage has been reached for the list of known environments.
+     */
+    stage?: ProgressReportStage;
+}
+
 export interface IDiscoveryAPI {
     /**
-     * Fires when the known list of environments starts refreshing, i.e when discovery starts or restarts.
+     * Tracks discovery progress for current list of known environments, i.e when it starts, finishes or any other relevant
+     * stage. Note the progress for a particular query is currently not tracked or reported, this only indicates progress of
+     * the entire collection.
      */
-    readonly onRefreshStart: Event<void>;
+    readonly onProgress: Event<ProgressNotificationEvent>;
     /**
      * Fires with details if the known list changes.
      */
@@ -183,7 +208,7 @@ export interface IDiscoveryAPI {
      * Resolves once environment list has finished refreshing, i.e all environments are
      * discovered. Carries `undefined` if there is no refresh currently going on.
      */
-    readonly refreshPromise: Promise<void> | undefined;
+    getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined;
     /**
      * Triggers a new refresh for query if there isn't any already running.
      */

src/client/pythonEnvironments/base/locatorUtils.ts

@@ -6,7 +6,14 @@ import { createDeferred } from '../../common/utils/async';
 import { getURIFilter } from '../../common/utils/misc';
 import { traceVerbose } from '../../logging';
 import { PythonEnvInfo } from './info';
-import { IPythonEnvsIterator, PythonEnvUpdatedEvent, PythonLocatorQuery } from './locator';
+import {
+    IPythonEnvsIterator,
+    isProgressEvent,
+    ProgressNotificationEvent,
+    ProgressReportStage,
+    PythonEnvUpdatedEvent,
+    PythonLocatorQuery,
+} from './locator';
 
 /**
  * Create a filter function to match the given query.
@@ -71,8 +78,11 @@ export async function getEnvs<I = PythonEnvInfo>(iterator: IPythonEnvsIterator<I
     if (iterator.onUpdated === undefined) {
         updatesDone.resolve();
     } else {
-        const listener = iterator.onUpdated((event: PythonEnvUpdatedEvent<I> | null) => {
-            if (event === null) {
+        const listener = iterator.onUpdated((event: PythonEnvUpdatedEvent<I> | ProgressNotificationEvent) => {
+            if (isProgressEvent(event)) {
+                if (event.stage !== ProgressReportStage.discoveryFinished) {
+                    return;
+                }
                 updatesDone.resolve();
                 listener.dispose();
             } else {

src/client/pythonEnvironments/base/locators.ts

@@ -4,7 +4,15 @@
 import { chain } from '../../common/utils/async';
 import { Disposables } from '../../common/utils/resourceLifecycle';
 import { PythonEnvInfo } from './info';
-import { ILocator, IPythonEnvsIterator, PythonEnvUpdatedEvent, PythonLocatorQuery } from './locator';
+import {
+    ILocator,
+    IPythonEnvsIterator,
+    isProgressEvent,
+    ProgressNotificationEvent,
+    ProgressReportStage,
+    PythonEnvUpdatedEvent,
+    PythonLocatorQuery,
+} from './locator';
 import { PythonEnvsWatchers } from './watchers';
 
 /**
@@ -19,17 +27,21 @@ export function combineIterators<I>(iterators: IPythonEnvsIterator<I>[]): IPytho
     }
 
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    result.onUpdated = (handleEvent: (e: PythonEnvUpdatedEvent<I> | null) => any) => {
+    result.onUpdated = (handleEvent: (e: PythonEnvUpdatedEvent<I> | ProgressNotificationEvent) => any) => {
         const disposables = new Disposables();
         let numActive = events.length;
         events.forEach((event) => {
-            const disposable = event!((e: PythonEnvUpdatedEvent<I> | null) => {
+            const disposable = event!((e: PythonEnvUpdatedEvent<I> | ProgressNotificationEvent) => {
                 // NOSONAR
-                if (e === null) {
-                    numActive -= 1;
-                    if (numActive === 0) {
-                        // All the sub-events are done so we're done.
-                        handleEvent(null);
+                if (isProgressEvent(e)) {
+                    if (e.stage === ProgressReportStage.discoveryFinished) {
+                        numActive -= 1;
+                        if (numActive === 0) {
+                            // All the sub-events are done so we're done.
+                            handleEvent({ stage: ProgressReportStage.discoveryFinished });
+                        }
+                    } else {
+                        handleEvent({ stage: e.stage });
                     }
                 } else {
                     handleEvent(e);