Show interpreter picker while environment discovery runs

src/common/pickers/environments.ts

@@ -8,6 +8,7 @@ import { sendTelemetryEvent } from '../telemetry/sender';
 import { isWindows } from '../utils/platformUtils';
 import { handlePythonPath } from '../utils/pythonPath';
 import {
+    QuickPickWithButtonsOptions,
     showErrorMessage,
     showInputBox,
     showOpenDialog,
@@ -47,6 +48,9 @@ interface EnvironmentPickOptions {
     showBackButton?: boolean;
     projects: PythonProject[];
 }
+
+type EnvironmentPickItem = QuickPickItem | (QuickPickItem & { result: PythonEnvironment });
+
 async function browseForPython(
     managers: InternalEnvironmentManager[],
     projectEnvManagers: InternalEnvironmentManager[],
@@ -144,15 +148,17 @@ async function createEnvironment(
 }
 
 async function pickEnvironmentImpl(
-    items: (QuickPickItem | (QuickPickItem & { result: PythonEnvironment }))[],
+    items: EnvironmentPickItem[],
     managers: InternalEnvironmentManager[],
     projectEnvManagers: InternalEnvironmentManager[],
     options: EnvironmentPickOptions,
+    itemLoader?: QuickPickWithButtonsOptions<EnvironmentPickItem>['itemLoader'],
 ): Promise<PythonEnvironment | undefined> {
     const selected = await showQuickPickWithButtons(items, {
         placeHolder: Pickers.Environments.selectEnvironment,
         ignoreFocusOut: true,
         showBackButton: options?.showBackButton,
+        itemLoader,
     });
 
     if (selected && !Array.isArray(selected)) {
@@ -177,7 +183,7 @@ export async function pickEnvironment(
     projectEnvManagers: InternalEnvironmentManager[],
     options: EnvironmentPickOptions,
 ): Promise<PythonEnvironment | undefined> {
-    const items: (QuickPickItem | (QuickPickItem & { result: PythonEnvironment }))[] = [
+    const items: EnvironmentPickItem[] = [
         {
             label: Interpreter.enterInterpreterPath,
             iconPath: new ThemeIcon('edit'),
@@ -222,30 +228,45 @@ export async function pickEnvironment(
         );
     }
 
-    for (const manager of managers) {
-        items.push({
-            label: manager.displayName,
-            kind: QuickPickItemKind.Separator,
-        });
-        const envs = await manager.getEnvironments('all');
-        items.push(
-            ...envs.map((e) => {
-                const pathDescription = e.displayPath;
-                const description =
-                    e.description && e.description.trim() ? `${e.description} (${pathDescription})` : pathDescription;
+    const managerItems: (EnvironmentPickItem[] | undefined)[] = [];
+    const itemLoader: QuickPickWithButtonsOptions<EnvironmentPickItem>['itemLoader'] =
+        managers.length > 0
+            ? async (updateItems) => {
+                  await Promise.all(
+                      managers.map(async (manager, index) => {
+                          try {
+                              const envs = await manager.getEnvironments('all');
+                              managerItems[index] = [
+                                  {
+                                      label: manager.displayName,
+                                      kind: QuickPickItemKind.Separator,
+                                  },
+                                  ...envs.map((e) => {
+                                      const pathDescription = e.displayPath;
+                                      const description =
+                                          e.description && e.description.trim()
+                                              ? `${e.description} (${pathDescription})`
+                                              : pathDescription;
 
-                return {
-                    label: e.displayName ?? e.name,
-                    description: description,
-                    result: e,
-                    manager: manager,
-                    iconPath: getIconPath(e.iconPath),
-                };
-            }),
-        );
-    }
+                                      return {
+                                          label: e.displayName ?? e.name,
+                                          description: description,
+                                          result: e,
+                                          manager: manager,
+                                          iconPath: getIconPath(e.iconPath),
+                                      };
+                                  }),
+                              ];
+                              updateItems([...items, ...managerItems.flatMap((current) => current ?? [])]);
+                          } catch (ex) {
+                              traceError(`Failed to get environments from ${manager.id}`, ex);
+                          }
+                      }),
+                  );
+              }
+            : undefined;
 
-    return pickEnvironmentImpl(items, managers, projectEnvManagers, options);
+    return pickEnvironmentImpl(items, managers, projectEnvManagers, options, itemLoader);
 }
 
 export async function pickEnvironmentFrom(environments: PythonEnvironment[]): Promise<PythonEnvironment | undefined> {

src/common/window.apis.ts

@@ -32,6 +32,7 @@ import {
     window,
     WindowState,
 } from 'vscode';
+import { traceError } from './logging';
 import { createDeferred } from './utils/deferred';
 
 export function createStatusBarItem(id: string, alignment?: StatusBarAlignment, priority?: number): StatusBarItem {
@@ -152,6 +153,13 @@ export interface QuickPickButtonEvent<T extends QuickPickItem> {
     readonly button: QuickInputButton;
 }
 
+export interface QuickPickWithButtonsOptions<T extends QuickPickItem> extends QuickPickOptions {
+    showBackButton?: boolean;
+    buttons?: QuickInputButton[];
+    selected?: T[];
+    itemLoader?: (updateItems: (items: readonly T[]) => void) => Promise<void>;
+}
+
 export function showQuickPick<T extends QuickPickItem>(
     items: readonly T[] | Thenable<readonly T[]>,
     options?: QuickPickOptions,
@@ -175,7 +183,7 @@ export function withProgress<R>(
 
 export async function showQuickPickWithButtons<T extends QuickPickItem>(
     items: readonly T[],
-    options?: QuickPickOptions & { showBackButton?: boolean; buttons?: QuickInputButton[]; selected?: T[] },
+    options?: QuickPickWithButtonsOptions<T>,
     token?: CancellationToken,
     itemButtonHandler?: (e: QuickPickItemButtonEvent<T>) => void,
 ): Promise<T | T[] | undefined> {
@@ -191,6 +199,8 @@ export async function showQuickPickWithButtons<T extends QuickPickItem>(
     quickPick.placeholder = options?.placeHolder;
     quickPick.title = options?.title;
     quickPick.selectedItems = options?.selected ?? [];
+    const itemLoader = options?.itemLoader;
+    quickPick.busy = itemLoader !== undefined;
 
     if (options?.showBackButton) {
         quickPick.buttons = [QuickInputButtons.Back];
@@ -241,6 +251,35 @@ export async function showQuickPickWithButtons<T extends QuickPickItem>(
     }
     quickPick.show();
 
+    if (itemLoader) {
+        void Promise.resolve()
+            .then(() =>
+                itemLoader((updatedItems) => {
+                    if (deferred.completed) {
+                        return;
+                    }
+                    const activeItems = quickPick.activeItems;
+                    const selectedItems = quickPick.selectedItems;
+                    quickPick.keepScrollPosition = true;
+                    quickPick.items = updatedItems;
+                    const retainedSelectedItems = selectedItems.filter((item) => updatedItems.includes(item));
+                    const retainedActiveItems = activeItems.filter((item) => updatedItems.includes(item));
+                    if (retainedSelectedItems.length > 0) {
+                        quickPick.selectedItems = retainedSelectedItems;
+                    }
+                    if (retainedActiveItems.length > 0) {
+                        quickPick.activeItems = retainedActiveItems;
+                    }
+                }),
+            )
+            .catch((ex) => traceError('Failed to update quick pick items', ex))
+            .finally(() => {
+                if (!deferred.completed) {
+                    quickPick.busy = false;
+                }
+            });
+    }
+
     try {
         return await deferred.promise;
     } finally {

src/test/common/apiWrappers.unit.test.ts

@@ -1,8 +1,20 @@
 import assert from 'node:assert';
 import { reset, verify, when } from 'ts-mockito';
-import { Disposable, MessageItem, MessageOptions, TextDocument, WindowState } from 'vscode';
+import {
+    Disposable,
+    EventEmitter,
+    MessageItem,
+    MessageOptions,
+    QuickInputButton,
+    QuickPick,
+    QuickPickItem,
+    QuickPickItemButtonEvent,
+    TextDocument,
+    WindowState,
+} from 'vscode';
 import { executeCommand, registerCommand } from '../../common/command.api';
-import { onDidChangeWindowState, showWarningMessage } from '../../common/window.apis';
+import { createDeferred } from '../../common/utils/deferred';
+import { onDidChangeWindowState, showQuickPickWithButtons, showWarningMessage } from '../../common/window.apis';
 import { onDidOpenTextDocument } from '../../common/workspace.apis';
 import { mockedVSCodeNamespaces } from '../unittests';
 
@@ -11,6 +23,67 @@ suite('Typed API wrappers', () => {
     const window = mockedVSCodeNamespaces.window!;
     const workspace = mockedVSCodeNamespaces.workspace!;
 
+    function createFakeQuickPick<T extends QuickPickItem>() {
+        const didAccept = new EventEmitter<void>();
+        const didHide = new EventEmitter<void>();
+        const didTriggerButton = new EventEmitter<QuickInputButton>();
+        const didTriggerItemButton = new EventEmitter<QuickPickItemButtonEvent<T>>();
+        let items: readonly T[] = [];
+        let activeItems: readonly T[] = [];
+        let selectedItems: readonly T[] = [];
+        let activeItemsSetCount = 0;
+        let showCount = 0;
+        let disposeCount = 0;
+        const quickPick = {
+            get items() {
+                return items;
+            },
+            set items(value: readonly T[]) {
+                items = value;
+                activeItems = [];
+                selectedItems = [];
+            },
+            get activeItems() {
+                return activeItems;
+            },
+            set activeItems(value: readonly T[]) {
+                activeItems = value;
+                activeItemsSetCount += 1;
+            },
+            get selectedItems() {
+                return selectedItems;
+            },
+            set selectedItems(value: readonly T[]) {
+                selectedItems = value;
+            },
+            buttons: [],
+            show: () => {
+                showCount += 1;
+            },
+            hide: () => didHide.fire(),
+            dispose: () => {
+                disposeCount += 1;
+            },
+            onDidAccept: didAccept.event,
+            onDidHide: didHide.event,
+            onDidTriggerButton: didTriggerButton.event,
+            onDidTriggerItemButton: didTriggerItemButton.event,
+        } as unknown as QuickPick<T>;
+        return {
+            quickPick,
+            hide: () => didHide.fire(),
+            get showCount() {
+                return showCount;
+            },
+            get disposeCount() {
+                return disposeCount;
+            },
+            get activeItemsSetCount() {
+                return activeItemsSetCount;
+            },
+        };
+    }
+
     teardown(() => {
         reset(commands);
         reset(window);
@@ -115,4 +188,93 @@ suite('Typed API wrappers', () => {
         assert.strictEqual(selected, action);
         assert.strictEqual(selected?.id, 1);
     });
+
+    test('updates quick pick items without losing focus while loading', async () => {
+        const fake = createFakeQuickPick<QuickPickItem>();
+        const initial = { label: 'initial' };
+        const discovered = { label: 'discovered' };
+        const release = createDeferred<void>();
+        const loaderStarted = createDeferred<void>();
+        const loaderFinished = createDeferred<void>();
+        when(window.createQuickPick<QuickPickItem>()).thenReturn(fake.quickPick);
+
+        const selection = showQuickPickWithButtons([initial], {
+            itemLoader: async (updateItems) => {
+                loaderStarted.resolve();
+                await release.promise;
+                updateItems([initial, discovered]);
+                loaderFinished.resolve();
+            },
+        });
+        await loaderStarted.promise;
+
+        fake.quickPick.activeItems = [initial];
+        fake.quickPick.selectedItems = [initial];
+        assert.strictEqual(fake.quickPick.busy, true);
+        assert.strictEqual(fake.showCount, 1);
+
+        release.resolve();
+        await loaderFinished.promise;
+        await new Promise<void>((resolve) => setImmediate(resolve));
+
+        assert.deepStrictEqual(fake.quickPick.items, [initial, discovered]);
+        assert.deepStrictEqual(fake.quickPick.activeItems, [initial]);
+        assert.deepStrictEqual(fake.quickPick.selectedItems, [initial]);
+        assert.strictEqual(fake.activeItemsSetCount, 2);
+        assert.strictEqual(fake.quickPick.keepScrollPosition, true);
+        assert.strictEqual(fake.quickPick.busy, false);
+
+        fake.hide();
+        assert.strictEqual(await selection, undefined);
+    });
+
+    test('ignores item updates after the quick pick is dismissed', async () => {
+        const fake = createFakeQuickPick<QuickPickItem>();
+        const initial = { label: 'initial' };
+        const release = createDeferred<void>();
+        const loaderStarted = createDeferred<void>();
+        const loaderFinished = createDeferred<void>();
+        when(window.createQuickPick<QuickPickItem>()).thenReturn(fake.quickPick);
+
+        const selection = showQuickPickWithButtons([initial], {
+            itemLoader: async (updateItems) => {
+                loaderStarted.resolve();
+                await release.promise;
+                updateItems([{ label: 'late result' }]);
+                loaderFinished.resolve();
+            },
+        });
+        await loaderStarted.promise;
+
+        fake.hide();
+        assert.strictEqual(await selection, undefined);
+        assert.strictEqual(fake.disposeCount, 1);
+
+        release.resolve();
+        await loaderFinished.promise;
+        await new Promise<void>((resolve) => setImmediate(resolve));
+
+        assert.deepStrictEqual(fake.quickPick.items, [initial]);
+        assert.strictEqual(fake.showCount, 1);
+    });
+
+    test('does not clear focus when items load before the initial focus is established', async () => {
+        const fake = createFakeQuickPick<QuickPickItem>();
+        const loaderFinished = createDeferred<void>();
+        when(window.createQuickPick<QuickPickItem>()).thenReturn(fake.quickPick);
+
+        const selection = showQuickPickWithButtons([{ label: 'initial' }], {
+            itemLoader: async (updateItems) => {
+                updateItems([{ label: 'initial' }, { label: 'discovered' }]);
+                loaderFinished.resolve();
+            },
+        });
+        await loaderFinished.promise;
+        await new Promise<void>((resolve) => setImmediate(resolve));
+
+        assert.strictEqual(fake.activeItemsSetCount, 0);
+
+        fake.hide();
+        assert.strictEqual(await selection, undefined);
+    });
 });