Direct users to the Jupyter extension when using Run in Interactive window (#21072)

Closes #20576

Author: Kartik Raj

package.json

@@ -374,6 +374,11 @@
                     "light": "resources/light/repl.svg"
                 },
                 "title": "%python.command.python.viewOutput.title%"
+            },
+            {
+                "category": "Python",
+                "command": "python.installJupyter",
+                "title": "%python.command.python.installJupyter.title%"
             }
         ],
         "configuration": {
@@ -1705,13 +1710,25 @@
                 {
                     "submenu": "python.run",
                     "group": "Python",
-                    "when": "editorLangId == python && !virtualWorkspace && shellExecutionSupported"
+                    "when": "editorLangId == python && !virtualWorkspace && shellExecutionSupported && isWorkspaceTrusted"
                 },
                 {
                     "command": "python.sortImports",
                     "group": "Refactor",
                     "title": "%python.command.python.sortImports.title%",
                     "when": "editorLangId == python && !notebookEditorFocused && !virtualWorkspace && shellExecutionSupported"
+                },
+                {
+                    "submenu": "python.runFileInteractive",
+                    "group": "Jupyter2",
+                    "when": "editorLangId == python && !virtualWorkspace && shellExecutionSupported && !isJupyterInstalled && isWorkspaceTrusted"
+                }
+            ],
+            "python.runFileInteractive": [
+                {
+                    "command": "python.installJupyter",
+                    "group": "Jupyter2",
+                    "when": "resourceLangId == python && !virtualWorkspace && shellExecutionSupported"
                 }
             ],
             "python.run": [
@@ -1779,6 +1796,10 @@
                 "id": "python.run",
                 "label": "%python.editor.context.submenu.runPython%",
                 "icon": "$(play)"
+            },
+            {
+                "id": "python.runFileInteractive",
+                "label": "%python.editor.context.submenu.runPythonInteractive%"
             }
         ],
         "viewsWelcome": [

package.nls.json

@@ -10,6 +10,7 @@
     "python.command.python.setInterpreter.title": "Select Interpreter",
     "python.command.python.clearWorkspaceInterpreter.title": "Clear Workspace Interpreter Setting",
     "python.command.python.viewOutput.title": "Show Output",
+    "python.command.python.installJupyter.title": "Install the Jupyter extension",
     "python.command.python.viewLanguageServerOutput.title": "Show Language Server Output",
     "python.command.python.configureTests.title": "Configure Tests",
     "python.command.testing.rerunFailedTests.title": "Rerun Failed Tests",
@@ -26,6 +27,7 @@
     "python.command.python.refreshTensorBoard.title": "Refresh TensorBoard",
     "python.menu.createNewFile.title": "Python File",
     "python.editor.context.submenu.runPython": "Run Python",
+    "python.editor.context.submenu.runPythonInteractive": "Run in Interactive window",
     "python.activeStateToolPath.description": "Path to the State Tool executable for ActiveState runtimes (version 0.36+).",
     "python.autoComplete.extraPaths.description": "List of paths to libraries and the like that need to be imported by auto complete engine. E.g. when using Google App SDK, the paths are not in system path, hence need to be added into this list.",
     "python.condaPath.description": "Path to the conda executable to use for activation (version 4.4+).",

src/client/common/application/commands.ts

@@ -17,6 +17,7 @@ export type CommandsWithoutArgs = keyof ICommandNameWithoutArgumentTypeMapping;
  */
 interface ICommandNameWithoutArgumentTypeMapping {
     [Commands.InstallPythonOnMac]: [];
+    [Commands.InstallJupyter]: [];
     [Commands.InstallPythonOnLinux]: [];
     [Commands.InstallPython]: [];
     [Commands.ClearWorkspaceInterpreter]: [];

src/client/common/application/contextKeys.ts

@@ -5,4 +5,5 @@ export enum ExtensionContextKey {
     showInstallPythonTile = 'showInstallPythonTile',
     HasFailedTests = 'hasFailedTests',
     RefreshingTests = 'refreshingTests',
+    IsJupyterInstalled = 'isJupyterInstalled',
 }

src/client/common/constants.ts

@@ -46,6 +46,7 @@ export namespace Commands {
     export const Exec_Selection_In_Django_Shell = 'python.execSelectionInDjangoShell';
     export const Exec_Selection_In_Terminal = 'python.execSelectionInTerminal';
     export const GetSelectedInterpreterPath = 'python.interpreterPath';
+    export const InstallJupyter = 'python.installJupyter';
     export const InstallPython = 'python.installPython';
     export const InstallPythonOnLinux = 'python.installPythonOnLinux';
     export const InstallPythonOnMac = 'python.installPythonOnMac';

src/client/common/serviceRegistry.ts

@@ -90,6 +90,7 @@ import { IMultiStepInputFactory, MultiStepInputFactory } from './utils/multiStep
 import { Random } from './utils/random';
 import { ContextKeyManager } from './application/contextKeyManager';
 import { CreatePythonFileCommandHandler } from './application/commands/createPythonFile';
+import { RequireJupyterPrompt } from '../jupyter/requireJupyterPrompt';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingletonInstance<boolean>(IsWindows, IS_WINDOWS);
@@ -110,6 +111,10 @@ export function registerTypes(serviceManager: IServiceManager): void {
         IJupyterExtensionDependencyManager,
         JupyterExtensionDependencyManager,
     );
+    serviceManager.addSingleton<IExtensionSingleActivationService>(
+        IExtensionSingleActivationService,
+        RequireJupyterPrompt,
+    );
     serviceManager.addSingleton<IExtensionSingleActivationService>(
         IExtensionSingleActivationService,
         CreatePythonFileCommandHandler,

src/client/common/utils/localize.ts

@@ -188,6 +188,9 @@ export namespace LanguageService {
     );
 }
 export namespace Interpreters {
+    export const requireJupyter = l10n.t(
+        'Running in Interactive window requires Jupyter Extension. Would you like to install it? [Learn more](https://aka.ms/pythonJupyterSupport).',
+    );
     export const installingPython = l10n.t('Installing Python into Environment...');
     export const discovering = l10n.t('Discovering Python Interpreters');
     export const refreshing = l10n.t('Refreshing Python Interpreters');

src/client/jupyter/jupyterIntegration.ts

@@ -8,7 +8,7 @@ import { inject, injectable, named } from 'inversify';
 import { dirname } from 'path';
 import { CancellationToken, Event, Extension, Memento, Uri } from 'vscode';
 import type { SemVer } from 'semver';
-import { IWorkspaceService } from '../common/application/types';
+import { IContextKeyManager, IWorkspaceService } from '../common/application/types';
 import { JUPYTER_EXTENSION_ID, PYLANCE_EXTENSION_ID } from '../common/constants';
 import { InterpreterUri, ModuleInstallFlags } from '../common/installer/types';
 import {
@@ -35,6 +35,7 @@ import {
 import { PythonEnvironment } from '../pythonEnvironments/info';
 import { IDataViewerDataProvider, IJupyterUriProvider } from './types';
 import { PylanceApi } from '../activation/node/pylanceApi';
+import { ExtensionContextKey } from '../common/application/contextKeys';
 /**
  * This allows Python extension to update Product enum without breaking Jupyter.
  * I.e. we have a strict contract, else using numbers (in enums) is bound to break across products.
@@ -201,9 +202,11 @@ export class JupyterExtensionIntegration {
         @inject(IComponentAdapter) private pyenvs: IComponentAdapter,
         @inject(IWorkspaceService) private workspaceService: IWorkspaceService,
         @inject(ICondaService) private readonly condaService: ICondaService,
+        @inject(IContextKeyManager) private readonly contextManager: IContextKeyManager,
     ) {}
 
     public registerApi(jupyterExtensionApi: JupyterExtensionApi): JupyterExtensionApi | undefined {
+        this.contextManager.setContext(ExtensionContextKey.IsJupyterInstalled, true);
         if (!this.workspaceService.isTrusted) {
             this.workspaceService.onDidGrantWorkspaceTrust(() => this.registerApi(jupyterExtensionApi));
             return undefined;

src/client/jupyter/requireJupyterPrompt.ts

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { IExtensionSingleActivationService } from '../activation/types';
+import { IApplicationShell, ICommandManager } from '../common/application/types';
+import { Common, Interpreters } from '../common/utils/localize';
+import { Commands, JUPYTER_EXTENSION_ID } from '../common/constants';
+import { IDisposable, IDisposableRegistry } from '../common/types';
+import { sendTelemetryEvent } from '../telemetry';
+import { EventName } from '../telemetry/constants';
+
+@injectable()
+export class RequireJupyterPrompt implements IExtensionSingleActivationService {
+    public readonly supportedWorkspaceTypes = { untrustedWorkspace: false, virtualWorkspace: true };
+
+    constructor(
+        @inject(IApplicationShell) private readonly appShell: IApplicationShell,
+        @inject(ICommandManager) private readonly commandManager: ICommandManager,
+        @inject(IDisposableRegistry) private readonly disposables: IDisposable[],
+    ) {}
+
+    public async activate(): Promise<void> {
+        this.disposables.push(this.commandManager.registerCommand(Commands.InstallJupyter, () => this._showPrompt()));
+    }
+
+    public async _showPrompt(): Promise<void> {
+        const prompts = [Common.bannerLabelYes, Common.bannerLabelNo];
+        const telemetrySelections: ['Yes', 'No'] = ['Yes', 'No'];
+        const selection = await this.appShell.showInformationMessage(Interpreters.requireJupyter, ...prompts);
+        sendTelemetryEvent(EventName.REQUIRE_JUPYTER_PROMPT, undefined, {
+            selection: selection ? telemetrySelections[prompts.indexOf(selection)] : undefined,
+        });
+        if (!selection) {
+            return;
+        }
+        if (selection === prompts[0]) {
+            await this.commandManager.executeCommand(
+                'workbench.extensions.installExtension',
+                JUPYTER_EXTENSION_ID,
+                undefined,
+            );
+        }
+    }
+}

src/client/telemetry/constants.ts

@@ -30,6 +30,7 @@ export enum EventName {
     PYTHON_INTERPRETER_ACTIVATE_ENVIRONMENT_PROMPT = 'PYTHON_INTERPRETER_ACTIVATE_ENVIRONMENT_PROMPT',
     PYTHON_NOT_INSTALLED_PROMPT = 'PYTHON_NOT_INSTALLED_PROMPT',
     CONDA_INHERIT_ENV_PROMPT = 'CONDA_INHERIT_ENV_PROMPT',
+    REQUIRE_JUPYTER_PROMPT = 'REQUIRE_JUPYTER_PROMPT',
     ACTIVATED_CONDA_ENV_LAUNCH = 'ACTIVATED_CONDA_ENV_LAUNCH',
     ENVFILE_VARIABLE_SUBSTITUTION = 'ENVFILE_VARIABLE_SUBSTITUTION',
     ENVFILE_WORKSPACE = 'ENVFILE_WORKSPACE',

src/client/telemetry/index.ts

@@ -1315,6 +1315,22 @@ export interface IEventNamePropertyMapping {
          */
         selection: 'Allow' | 'Close' | undefined;
     };
+    /**
+     * Telemetry event sent with details when user attempts to run in interactive window when Jupyter is not installed.
+     */
+    /* __GDPR__
+       "conda_inherit_env_prompt" : {
+          "selection" : { "classification": "SystemMetaData", "purpose": "FeatureInsight", "owner": "karrtikr" }
+       }
+     */
+    [EventName.REQUIRE_JUPYTER_PROMPT]: {
+        /**
+         * `Yes` When 'Yes' option is selected
+         * `No` When 'No' option is selected
+         * `undefined` When 'x' is selected
+         */
+        selection: 'Yes' | 'No' | undefined;
+    };
     /**
      * Telemetry event sent with details when user clicks the prompt with the following message:
      *

src/test/activation/node/lspInteractiveWindowMiddlewareAddon.unit.test.ts

@@ -19,7 +19,7 @@ import {
 } from '../../../client/interpreter/contracts';
 import { IInterpreterSelector } from '../../../client/interpreter/configuration/types';
 import { IEnvironmentActivationService } from '../../../client/interpreter/activation/types';
-import { IWorkspaceService } from '../../../client/common/application/types';
+import { IContextKeyManager, IWorkspaceService } from '../../../client/common/application/types';
 import { MockMemento } from '../../mocks/mementos';
 
 suite('Pylance Language Server - Interactive Window LSP Notebooks', () => {
@@ -41,6 +41,7 @@ suite('Pylance Language Server - Interactive Window LSP Notebooks', () => {
             mock<IComponentAdapter>(),
             mock<IWorkspaceService>(),
             mock<ICondaService>(),
+            mock<IContextKeyManager>(),
         );
         jupyterApi.registerGetNotebookUriForTextDocumentUriFunction(getNotebookUriFunction);
     });

src/test/jupyter/requireJupyterPrompt.unit.test.ts

@@ -0,0 +1,60 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { mock, instance, verify, anything, when } from 'ts-mockito';
+import { IApplicationShell, ICommandManager } from '../../client/common/application/types';
+import { Commands, JUPYTER_EXTENSION_ID } from '../../client/common/constants';
+import { IDisposableRegistry } from '../../client/common/types';
+import { Common, Interpreters } from '../../client/common/utils/localize';
+import { RequireJupyterPrompt } from '../../client/jupyter/requireJupyterPrompt';
+
+suite('RequireJupyterPrompt Unit Tests', () => {
+    let requireJupyterPrompt: RequireJupyterPrompt;
+    let appShell: IApplicationShell;
+    let commandManager: ICommandManager;
+    let disposables: IDisposableRegistry;
+
+    setup(() => {
+        appShell = mock<IApplicationShell>();
+        commandManager = mock<ICommandManager>();
+        disposables = mock<IDisposableRegistry>();
+
+        requireJupyterPrompt = new RequireJupyterPrompt(
+            instance(appShell),
+            instance(commandManager),
+            instance(disposables),
+        );
+    });
+
+    test('Activation registers command', async () => {
+        await requireJupyterPrompt.activate();
+
+        verify(commandManager.registerCommand(Commands.InstallJupyter, anything())).once();
+    });
+
+    test('Show prompt with Yes selection installs Jupyter extension', async () => {
+        when(
+            appShell.showInformationMessage(Interpreters.requireJupyter, Common.bannerLabelYes, Common.bannerLabelNo),
+        ).thenReturn(Promise.resolve(Common.bannerLabelYes));
+
+        await requireJupyterPrompt.activate();
+        await requireJupyterPrompt._showPrompt();
+
+        verify(
+            commandManager.executeCommand('workbench.extensions.installExtension', JUPYTER_EXTENSION_ID, undefined),
+        ).once();
+    });
+
+    test('Show prompt with No selection does not install Jupyter extension', async () => {
+        when(
+            appShell.showInformationMessage(Interpreters.requireJupyter, Common.bannerLabelYes, Common.bannerLabelNo),
+        ).thenReturn(Promise.resolve(Common.bannerLabelNo));
+
+        await requireJupyterPrompt.activate();
+        await requireJupyterPrompt._showPrompt();
+
+        verify(
+            commandManager.executeCommand('workbench.extensions.installExtension', JUPYTER_EXTENSION_ID, undefined),
+        ).never();
+    });
+});