Two step process to select a source and a kernel in the insiders kern…

…el picker (#11664)

* two step process to select a source and a kernel

* add news

Co-authored-by: Ian Huff <[email protected]>

news/1 Enhancements/11642.md

@@ -0,0 +1 @@
+Make insiders kernel picker now a two step selection process that picks a source and a controller.

package.nls.json

@@ -1132,5 +1132,7 @@
     "jupyter.configuration.jupyter.diagnostics.reservedPythonNames.exclude.markdownDescription": "Configure [glob patterns](https://code.visualstudio.com/docs/editor/codebasics#_advanced-search-options) for excluding files and folders from from being warned about [overriding Python packages](https://aka.ms/JupyterKernelStartFailureOverrideReservedName).",
     "DataScience.localKernelFinderDisplayName": "Local Kernels",
     "DataScience.universalRemoteKernelFinderDisplayName": "Remote - {0}",
-    "DataScience.remoteKernelFinderDisplayName": "Current Remote"
+    "DataScience.remoteKernelFinderDisplayName": "Current Remote",
+    "DataScience.kernelPickerSelectSourceTitle": "Select Kernel Source",
+    "DataScience.kernelPickerSelectKernelTitle": "Select Kernel"
 }

src/notebooks/controllers/commands/serverConnectionControllerCommands.ts

@@ -30,7 +30,7 @@ import { isLocalConnection } from '../../../kernels/types';
 import { IJupyterServerUriStorage, IServerConnectionType } from '../../../kernels/jupyter/types';
 import { DataScience } from '../../../platform/common/utils/localize';
 
-function groupBy<T>(data: ReadonlyArray<T>, compare: (a: T, b: T) => number): T[][] {
+export function groupBy<T>(data: ReadonlyArray<T>, compare: (a: T, b: T) => number): T[][] {
     const result: T[][] = [];
     let currentGroup: T[] | undefined = undefined;
     for (const element of data.slice(0).sort(compare)) {
@@ -44,7 +44,7 @@ function groupBy<T>(data: ReadonlyArray<T>, compare: (a: T, b: T) => number): T[
     return result;
 }
 
-function compareIgnoreCase(a: string, b: string) {
+export function compareIgnoreCase(a: string, b: string) {
     return a.localeCompare(b, undefined, { sensitivity: 'accent' });
 }
 

src/notebooks/controllers/kernelSource/notebookKernelSourceSelector.ts

@@ -4,35 +4,163 @@
 'use strict';
 
 import { inject, injectable } from 'inversify';
-import { NotebookDocument, QuickPickItem } from 'vscode';
+import { NotebookDocument, QuickPickItem, QuickPickItemKind } from 'vscode';
 import { IContributedKernelFinderInfo } from '../../../kernels/internalTypes';
 import { IKernelFinder } from '../../../kernels/types';
-import { IApplicationShell } from '../../../platform/common/application/types';
-import { INotebookKernelSourceSelector, INotebookKernelSourceTracker } from '../types';
+import { ICommandManager } from '../../../platform/common/application/types';
+import { InteractiveWindowView, JupyterNotebookView, JVSC_EXTENSION_ID } from '../../../platform/common/constants';
+import { DataScience } from '../../../platform/common/utils/localize';
+import {
+    IMultiStepInput,
+    IMultiStepInputFactory,
+    IQuickPickParameters
+} from '../../../platform/common/utils/multiStepInput';
+import { compareIgnoreCase, groupBy } from '../commands/serverConnectionControllerCommands';
+import {
+    IControllerRegistration,
+    INotebookKernelSourceSelector,
+    INotebookKernelSourceTracker,
+    IVSCodeNotebookController
+} from '../types';
 
 interface KernelFinderQuickPickItem extends QuickPickItem {
     kernelFinderInfo: IContributedKernelFinderInfo;
 }
 
+interface ControllerQuickPickItem extends QuickPickItem {
+    controller: IVSCodeNotebookController;
+}
+
+// The return type of our multistep selection process
+type MultiStepResult = { source?: IContributedKernelFinderInfo; controller?: IVSCodeNotebookController };
+
 // Provides the UI to select a Kernel Source for a given notebook document
 @injectable()
 export class NotebookKernelSourceSelector implements INotebookKernelSourceSelector {
     constructor(
-        @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
         @inject(INotebookKernelSourceTracker) private readonly kernelSourceTracker: INotebookKernelSourceTracker,
-        @inject(IKernelFinder) private readonly kernelFinder: IKernelFinder
+        @inject(IKernelFinder) private readonly kernelFinder: IKernelFinder,
+        @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
+        @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration,
+        @inject(ICommandManager) private readonly commandManager: ICommandManager
     ) {}
 
     public async selectKernelSource(notebook: NotebookDocument): Promise<void> {
+        // Reject if it's not our type
+        if (notebook.notebookType !== JupyterNotebookView && notebook.notebookType !== InteractiveWindowView) {
+            return;
+        }
+
+        const multiStep = this.multiStepFactory.create<MultiStepResult>();
+        const state: MultiStepResult = {};
+        await multiStep.run(this.getSource.bind(this, notebook.notebookType), state);
+
+        // If we got both parts of the equation, then perform the kernel source and kernel switch
+        if (state.source && state.controller) {
+            await this.applyResults(notebook, state);
+        }
+    }
+
+    // The first stage of the multistep to get source and kernel
+    private async getSource(
+        notebookType: typeof JupyterNotebookView | typeof InteractiveWindowView,
+        multiStep: IMultiStepInput<MultiStepResult>,
+        state: MultiStepResult
+    ) {
         const quickPickItems = this.kernelFinder.registered.map(this.toQuickPickItem);
-        const selectedItem = await this.applicationShell.showQuickPick(quickPickItems);
+        const selectedSource = await multiStep.showQuickPick<
+            KernelFinderQuickPickItem,
+            IQuickPickParameters<KernelFinderQuickPickItem>
+        >({ items: quickPickItems, placeholder: '', title: DataScience.kernelPickerSelectSourceTitle() });
+
+        if (selectedSource) {
+            // Got a source, now get the kernel
+            state.source = selectedSource.kernelFinderInfo;
+            return this.getKernel.bind(this, notebookType);
+        }
+    }
 
-        // If we selected something persist that value
-        if (selectedItem) {
-            this.kernelSourceTracker.setKernelSourceForNotebook(notebook, selectedItem.kernelFinderInfo);
+    // Second stage of the multistep to pick a kernel
+    private async getKernel(
+        notebookType: typeof JupyterNotebookView | typeof InteractiveWindowView,
+        multiStep: IMultiStepInput<MultiStepResult>,
+        state: MultiStepResult
+    ) {
+        if (!state.source) {
+            return;
         }
+
+        const matchingControllers = this.getMatchingControllers(state.source, notebookType);
+
+        // Create controller items and group the by category
+        const controllerPickItems: ControllerQuickPickItem[] = matchingControllers.map((controller) => {
+            return {
+                label: controller.label,
+                detail: undefined,
+                description: controller.controller.description,
+                controller
+            };
+        });
+
+        const kernelsPerCategory = groupBy(controllerPickItems, (a, b) =>
+            compareIgnoreCase(a.controller.controller.kind || 'z', b.controller.controller.kind || 'z')
+        );
+
+        // Insert separators into the right spots in the list
+        const kindIndexes = new Map<string, number>();
+        const quickPickItems: (QuickPickItem | ControllerQuickPickItem)[] = [];
+
+        kernelsPerCategory.forEach((items) => {
+            const kind = items[0].controller.controller.kind || 'Other';
+            quickPickItems.push({
+                kind: QuickPickItemKind.Separator,
+                label: kind
+            });
+            quickPickItems.push(...items);
+            kindIndexes.set(kind, quickPickItems.length);
+        });
+
+        const result = await multiStep.showQuickPick<
+            ControllerQuickPickItem | QuickPickItem,
+            IQuickPickParameters<ControllerQuickPickItem | QuickPickItem>
+        >({
+            title: DataScience.kernelPickerSelectKernelTitle(),
+            items: quickPickItems,
+            matchOnDescription: true,
+            matchOnDetail: true,
+            placeholder: ''
+        });
+
+        if ('controller' in result) {
+            state.controller = result.controller;
+        }
+    }
+
+    // Get all registered controllers that match a specific finder
+    private getMatchingControllers(
+        kernelSource: IContributedKernelFinderInfo,
+        notebookType: typeof JupyterNotebookView | typeof InteractiveWindowView
+    ): IVSCodeNotebookController[] {
+        return this.controllerRegistration.registered.filter((controller) => {
+            const finder = this.kernelFinder.getFinderForConnection(controller.connection);
+            return finder?.id === kernelSource.id && controller.viewType === notebookType;
+        });
+    }
+
+    // If we completed the multistep with results, apply those results
+    private async applyResults(notebook: NotebookDocument, result: MultiStepResult) {
+        // First apply the kernel filter to this document
+        result.source && this.kernelSourceTracker.setKernelSourceForNotebook(notebook, result.source);
+
+        // Then select the kernel that we wanted
+        result.controller &&
+            (await this.commandManager.executeCommand('notebook.selectKernel', {
+                id: result.controller.id,
+                extension: JVSC_EXTENSION_ID
+            }));
     }
 
+    // Convert a kernel finder info in a quick pick item
     toQuickPickItem(kernelFinderInfo: IContributedKernelFinderInfo): KernelFinderQuickPickItem {
         return { kernelFinderInfo, label: kernelFinderInfo.displayName };
     }

src/platform/common/utils/localize.ts

@@ -1420,6 +1420,10 @@ export namespace DataScience {
         localize('DataScience.universalRemoteKernelFinderDisplayName', 'Remote - {0}');
     export const remoteKernelFinderDisplayName = () =>
         localize('DataScience.remoteKernelFinderDisplayName', 'Current Remote');
+    export const kernelPickerSelectSourceTitle = () =>
+        localize('DataScience.kernelPickerSelectSourceTitle', 'Select Kernel Source');
+    export const kernelPickerSelectKernelTitle = () =>
+        localize('DataScience.kernelPickerSelectKernelTitle', 'Select Kernel');
 }
 
 export namespace Deprecated {