Add Kernel selection picker for local and remote sessions (#8873)

* Support installing ipykernel as a product
* No need
* Linter fixes
* Address code review comments

Author: DonJayamanne

src/client/datascience/jupyter/kernels/kernelSelections.ts

@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { CancellationToken } from 'vscode';
+import * as localize from '../../../common/utils/localize';
+import { IInterpreterSelector } from '../../../interpreter/configuration/types';
+import { IJupyterKernel, IJupyterKernelSpec, IJupyterSessionManager } from '../../types';
+import { KernelService } from './kernelService';
+import { IKernelSelectionListProvider, IKernelSpecQuickPickItem } from './types';
+
+// Small classes, hence all put into one file.
+// tslint:disable: max-classes-per-file
+
+/**
+ * Given a kernel spec, this will return a quick pick item with appropriate display names and the like.
+ *
+ * @param {IJupyterKernelSpec} kernelSpec
+ * @returns {IKernelSpecQuickPickItem}
+ */
+function getQuickPickItemForKernelSpec(kernelSpec: IJupyterKernelSpec): IKernelSpecQuickPickItem {
+    return {
+        label: kernelSpec.display_name,
+        // tslint:disable-next-line: no-suspicious-comment
+        // TODO: Localize & fix as per spec.
+        description: '(kernel)',
+        selection: { kernelModel: undefined, kernelSpec: kernelSpec, interpreter: undefined }
+    };
+}
+
+/**
+ * Given an active kernel, this will return a quick pick item with appropriate display names and the like.
+ *
+ * @param {(IJupyterKernel & Partial<IJupyterKernelSpec>)} kernel
+ * @returns {IKernelSpecQuickPickItem}
+ */
+function getQuickPickItemForActiveKernel(kernel: IJupyterKernel & Partial<IJupyterKernelSpec>): IKernelSpecQuickPickItem {
+    return {
+        label: kernel.display_name || kernel.name || '',
+        description: localize.DataScience.jupyterSelectURIRunningDetailFormat().format(kernel.lastActivityTime.toLocaleString(), kernel.numberOfConnections.toString()),
+        selection: { kernelModel: kernel, kernelSpec: undefined, interpreter: undefined }
+    };
+}
+
+/**
+ * Provider for active kernel specs in a jupyter session.
+ *
+ * @export
+ * @class ActiveJupyterSessionKernelSelectionListProvider
+ * @implements {IKernelSelectionListProvider}
+ */
+export class ActiveJupyterSessionKernelSelectionListProvider implements IKernelSelectionListProvider {
+    constructor(private readonly sessionManager: IJupyterSessionManager) {}
+    public async getKernelSelections(_cancelToken?: CancellationToken | undefined): Promise<IKernelSpecQuickPickItem[]> {
+        const [activeKernels, kernelSpecs] = await Promise.all([this.sessionManager.getRunningKernels(), this.sessionManager.getKernelSpecs()]);
+        const items = activeKernels.map(item => {
+            const matchingSpec: Partial<IJupyterKernelSpec> = kernelSpecs.find(spec => spec.name === item.name) || {};
+            return {
+                ...item,
+                ...matchingSpec
+            };
+        });
+        return items.filter(item => item.display_name || item.name).map(getQuickPickItemForActiveKernel);
+    }
+}
+
+/**
+ * Provider for installed kernel specs (`python -m jupyter kernelspec list`).
+ *
+ * @export
+ * @class InstalledJupyterKernelSelectionListProvider
+ * @implements {IKernelSelectionListProvider}
+ */
+export class InstalledJupyterKernelSelectionListProvider implements IKernelSelectionListProvider {
+    constructor(private readonly kernelService: KernelService, private readonly sessionManager?: IJupyterSessionManager) {}
+    public async getKernelSelections(cancelToken?: CancellationToken | undefined): Promise<IKernelSpecQuickPickItem[]> {
+        const items = await this.kernelService.getKernelSpecs(this.sessionManager, cancelToken);
+        return items.map(getQuickPickItemForKernelSpec);
+    }
+}
+
+/**
+ * Provider for interpreters to be treated as kernel specs.
+ * I.e. return interpreters that are to be treated as kernel specs, and not yet installed as kernels.
+ *
+ * @export
+ * @class InterpreterKernelSelectionListProvider
+ * @implements {IKernelSelectionListProvider}
+ */
+export class InterpreterKernelSelectionListProvider implements IKernelSelectionListProvider {
+    constructor(private readonly interpreterSelector: IInterpreterSelector) {}
+    public async getKernelSelections(_cancelToken?: CancellationToken | undefined): Promise<IKernelSpecQuickPickItem[]> {
+        const items = await this.interpreterSelector.getSuggestions(undefined);
+        return items.map(item => {
+            return {
+                ...item,
+        // tslint:disable-next-line: no-suspicious-comment
+                // TODO: Localize & fix as per spec.
+                description: '(register and use interpreter as kernel)',
+                selection: { kernelModel: undefined, interpreter: item.interpreter, kernelSpec: undefined }
+            };
+        });
+    }
+}
+
+/**
+ * Provides a list of kernel specs for selection, for both local and remote sessions.
+ *
+ * @export
+ * @class KernelSelectionProviderFactory
+ */
+@injectable()
+export class KernelSelectionProvider {
+    constructor(@inject(KernelService) private readonly kernelService: KernelService, @inject(IInterpreterSelector) private readonly interpreterSelector: IInterpreterSelector) {}
+    /**
+     * Gets a selection of kernel specs from a remote session.
+     *
+     * @param {IJupyterSessionManager} sessionManager
+     * @param {CancellationToken} [cancelToken]
+     * @returns {Promise<IKernelSpecQuickPickItem[]>}
+     * @memberof KernelSelectionProvider
+     */
+    public async getKernelSelectionsForRemoteSession(sessionManager: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IKernelSpecQuickPickItem[]> {
+        return new ActiveJupyterSessionKernelSelectionListProvider(sessionManager).getKernelSelections(cancelToken);
+    }
+    /**
+     * Gets a selection of kernel specs for a local session.
+     *
+     * @param {IJupyterSessionManager} [sessionManager]
+     * @param {CancellationToken} [cancelToken]
+     * @returns {Promise<IKernelSelectionListProvider>}
+     * @memberof KernelSelectionProvider
+     */
+    public async getKernelSelectionsForLocalSession(sessionManager?: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IKernelSpecQuickPickItem[]> {
+        const activeKernelsPromise = sessionManager ? new ActiveJupyterSessionKernelSelectionListProvider(sessionManager).getKernelSelections(cancelToken) : Promise.resolve([]);
+        const jupyterKernelsPromise = new InstalledJupyterKernelSelectionListProvider(this.kernelService).getKernelSelections(cancelToken);
+        const interpretersPromise = new InterpreterKernelSelectionListProvider(this.interpreterSelector).getKernelSelections(cancelToken);
+        const [activeKernels, jupyterKernels, interprters] = await Promise.all([activeKernelsPromise, jupyterKernelsPromise, interpretersPromise]);
+        return [...jupyterKernels!, ...activeKernels!, ...interprters];
+    }
+}

src/client/datascience/jupyter/kernels/kernelSelector.ts

@@ -0,0 +1,82 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { CancellationToken } from 'vscode-jsonrpc';
+import { IApplicationShell } from '../../../common/application/types';
+import { Cancellation } from '../../../common/cancellation';
+import { traceInfo, traceWarning } from '../../../common/logger';
+import { IInstaller, InstallerResponse, Product } from '../../../common/types';
+import { PythonInterpreter } from '../../../interpreter/contracts';
+import { IJupyterKernelSpec, IJupyterSessionManager } from '../../types';
+import { KernelSelectionProvider } from './kernelSelections';
+import { KernelService } from './kernelService';
+
+@injectable()
+export class KernelSelector {
+    constructor(
+        @inject(KernelSelectionProvider) private readonly selectionProvider: KernelSelectionProvider,
+        @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
+        @inject(KernelService) private readonly kernelService: KernelService,
+        @inject(IInstaller) private readonly installer: IInstaller
+    ) {}
+    public async selectRemoteKernel(session: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IJupyterKernelSpec | undefined> {
+        const suggestions = this.selectionProvider.getKernelSelectionsForRemoteSession(session, cancelToken);
+        const selection = await this.applicationShell.showQuickPick(suggestions, undefined, cancelToken);
+        if (!selection) {
+            return;
+        }
+
+        if (selection.selection.kernelSpec) {
+            return selection.selection.kernelSpec;
+        }
+        // This is not possible (remote kernels selector can only display remote kernels).
+        throw new Error('Invalid Selection in kernel spec (somehow a local kernel/interpreter has been selected for a remote session!');
+    }
+
+    public async selectLocalKernel(session?: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IJupyterKernelSpec | undefined> {
+        const suggestions = this.selectionProvider.getKernelSelectionsForLocalSession(session, cancelToken);
+        const selection = await this.applicationShell.showQuickPick(suggestions, undefined, cancelToken);
+        if (!selection) {
+            return;
+        }
+
+        // Check if ipykernel is installed in this kernel.
+        const interpreter = selection.selection.interpreter;
+        if (interpreter) {
+            const isValid = await this.isSelectionValid(interpreter, cancelToken);
+            if (isValid) {
+                // Find the kernel associated with this interpter.
+                const kernelSpec = await this.kernelService.findMatchingKernelSpec(interpreter, session, cancelToken);
+                if (kernelSpec){
+                    traceInfo(`ipykernel installed in ${interpreter.path}, and matching found.`);
+                    return kernelSpec;
+                }
+                traceInfo(`ipykernel installed in ${interpreter.path}, no matching kernel found. Will register kernel.`);
+            }
+
+            // Try an install this interpreter as a kernel.
+            return this.kernelService.registerKernel(interpreter, cancelToken);
+        } else {
+            return selection.selection.kernelSpec;
+        }
+    }
+
+    private async isSelectionValid(interpreter: PythonInterpreter, cancelToken?: CancellationToken): Promise<boolean> {
+        // Is ipykernel installed in this environment.
+        if (await this.installer.isInstalled(Product.ipykernel, interpreter)) {
+            return true;
+        }
+        if (Cancellation.isCanceled(cancelToken)) {
+            return false;
+        }
+        const response = await this.installer.promptToInstall(Product.ipykernel, interpreter);
+        if (response === InstallerResponse.Installed) {
+            return true;
+        }
+        traceWarning(`Prompted to install ipykernel, however ipykernel not installed in the interpreter ${interpreter.path}`);
+        return false;
+    }
+}

src/client/datascience/jupyter/kernels/kernelService.ts

@@ -221,6 +221,22 @@ export class KernelService {
         traceInfo(`Kernel successfully registered for ${interpreter.path} with the name=${name} and spec can be found here ${kernel.specFile}`);
         return kernel;
     }
+    /**
+     * Gets a list of all kernel specs.
+     *
+     * @param {IJupyterSessionManager} [sessionManager]
+     * @param {CancellationToken} [cancelToken]
+     * @returns {Promise<IJupyterKernelSpec[]>}
+     * @memberof KernelService
+     */
+    public async getKernelSpecs(sessionManager?: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IJupyterKernelSpec[]> {
+        const enumerator = sessionManager ? sessionManager.getKernelSpecs() : this.enumerateSpecs(cancelToken);
+        if (Cancellation.isCanceled(cancelToken)) {
+            return [];
+        }
+        const specs = await enumerator;
+        return specs.filter(item => !!item);
+    }
     /**
      * Not all characters are allowed in a kernel name.
      * This method will generate a name for a kernel based on display name and path.
@@ -233,14 +249,6 @@ export class KernelService {
     private async generateKernelNameForIntepreter(interpreter: PythonInterpreter): Promise<string> {
         return `${interpreter.displayName || ''}_${await this.fileSystem.getFileHash(interpreter.path)}`.replace(/[^A-Za-z0-9]/g, '');
     }
-    private async getKernelSpecs(sessionManager?: IJupyterSessionManager, cancelToken?: CancellationToken): Promise<IJupyterKernelSpec[]> {
-        const enumerator = sessionManager ? sessionManager.getKernelSpecs() : this.enumerateSpecs(cancelToken);
-        if (Cancellation.isCanceled(cancelToken)) {
-            return [];
-        }
-        const specs = await enumerator;
-        return specs.filter(item => !!item);
-    }
     private hasSpecPathMatch = async (info: PythonInterpreter | undefined, cancelToken?: CancellationToken): Promise<boolean> => {
         if (info) {
             // Enumerate our specs

src/client/datascience/jupyter/kernels/types.ts

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { CancellationToken, QuickPickItem } from 'vscode';
+import { PythonInterpreter } from '../../../interpreter/contracts';
+import { IJupyterKernel, IJupyterKernelSpec } from '../../types';
+
+export interface IKernelSpecQuickPickItem extends QuickPickItem {
+    /**
+     * Whether a
+     * - Kernel spec (IJupyterKernelSpec)
+     * - Active kernel (IJupyterKernel) or
+     * - Interpreter has been selected.
+     * If interpreter is selected, then we might need to install this as a kernel to get the kernel spec.
+     *
+     * @type {({ kernelModel: IJupyterKernel; kernelSpec: IJupyterKernelSpec; interpreter: undefined }
+     *         | { kernelModel: undefined; kernelSpec: IJupyterKernelSpec; interpreter: undefined }
+     *         | { kernelModel: undefined; kernelSpec: undefined; interpreter: PythonInterpreter })}
+     * @memberof IKernelSpecQuickPickItem
+     */
+    selection:
+        | { kernelModel: IJupyterKernel & Partial<IJupyterKernelSpec>; kernelSpec: undefined; interpreter: undefined }
+        | { kernelModel: undefined; kernelSpec: IJupyterKernelSpec; interpreter: undefined }
+        | { kernelModel: undefined; kernelSpec: undefined; interpreter: PythonInterpreter };
+}
+
+export interface IKernelSelectionListProvider {
+    getKernelSelections(cancelToken?: CancellationToken): Promise<IKernelSpecQuickPickItem[]>;
+}

src/client/datascience/serviceRegistry.ts

@@ -38,6 +38,8 @@ import { JupyterPasswordConnect } from './jupyter/jupyterPasswordConnect';
 import { JupyterServerFactory } from './jupyter/jupyterServerFactory';
 import { JupyterSessionManagerFactory } from './jupyter/jupyterSessionManagerFactory';
 import { JupyterVariables } from './jupyter/jupyterVariables';
+import { KernelSelectionProvider } from './jupyter/kernels/kernelSelections';
+import { KernelSelector } from './jupyter/kernels/kernelSelector';
 import { PlotViewer } from './plotting/plotViewer';
 import { PlotViewerProvider } from './plotting/plotViewerProvider';
 import { StatusProvider } from './statusProvider';
@@ -121,4 +123,6 @@ export function registerTypes(serviceManager: IServiceManager) {
     serviceManager.addSingleton<IDebugLocationTracker>(IDebugLocationTracker, DebugLocationTrackerFactory);
     serviceManager.addSingleton<JupyterCommandFinder>(JupyterCommandFinder, JupyterCommandFinder);
     serviceManager.addSingleton<IExtensionSingleActivationService>(IExtensionSingleActivationService, Activation);
+    serviceManager.addSingleton<KernelSelector>(KernelSelector, KernelSelector);
+    serviceManager.addSingleton<KernelSelectionProvider>(KernelSelectionProvider, KernelSelectionProvider);
 }