Add get refresh promise options

Author: Kartik Raj

src/client/apiTypes.ts

@@ -5,7 +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 { ProgressNotificationEvent } from './pythonEnvironments/base/locator';
+import { GetRefreshEnvironmentsOptions, ProgressNotificationEvent } from './pythonEnvironments/base/locator';
 
 /*
  * Do not introduce any breaking changes to this API.
@@ -212,7 +212,7 @@ export interface IProposedExtensionAPI {
          * 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/pythonEnvironments/base/locator.ts

@@ -186,6 +186,10 @@ interface IResolver {
 
 export interface IResolvingLocator<I = PythonEnvInfo> extends IResolver, ILocator<I> {}
 
+export interface GetRefreshEnvironmentsOptions {
+    stage?: ProgressReportStage;
+}
+
 export interface IDiscoveryAPI {
     /**
      * Fires with details of the current discovery progress, i.e when it starts, finishes or any other relevant stage.

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

@@ -9,7 +9,6 @@ import { normalizePath } from '../../../common/externalDependencies';
 import { PythonEnvInfo } from '../../info';
 import {
     IDiscoveryAPI,
-    IPythonEnvsIterator,
     IResolvingLocator,
     isProgressEvent,
     ProgressNotificationEvent,
@@ -100,14 +99,12 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
 
     private startRefresh(query: (PythonLocatorQuery & { clearCache?: boolean }) | undefined): Promise<void> {
         const deferred = createDeferred<void>();
-
         if (query?.clearCache) {
             this.cache.clearCache();
         }
         // Ensure we set this before we trigger the promise to accurately track when a refresh has started.
         this.refreshPromises.set(query, deferred.promise);
-        const iterator = this.locator.iterEnvs(query);
-        const promise = this.addEnvsToCacheFromIterator(iterator);
+        const promise = this.addEnvsToCacheForQuery(query);
         return promise
             .then(async () => {
                 // Ensure we delete this before we resolve the promise to accurately track when a refresh finishes.
@@ -117,7 +114,8 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
             .catch((ex) => deferred.reject(ex));
     }
 
-    private async addEnvsToCacheFromIterator(iterator: IPythonEnvsIterator) {
+    private async addEnvsToCacheForQuery(query: PythonLocatorQuery | undefined) {
+        const iterator = this.locator.iterEnvs(query);
         const seen: PythonEnvInfo[] = [];
         const state = {
             done: false,
@@ -128,12 +126,12 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
         if (iterator.onUpdated !== undefined) {
             const listener = iterator.onUpdated(async (event) => {
                 if (isProgressEvent(event)) {
-                    if (event.stage !== ProgressReportStage.discoveryFinished) {
+                    if (event.stage === ProgressReportStage.discoveryFinished) {
+                        state.done = true;
+                        listener.dispose();
+                    } else if (event.stage === ProgressReportStage.allPathsDiscovered && !query) {
                         this.progress.fire(event);
-                        return;
                     }
-                    state.done = true;
-                    listener.dispose();
                 } else {
                     state.pending += 1;
                     this.cache.updateEnv(seen[event.index], event.update);