Diagnostic provider for '!pip install' and `!conda install'

package.nls.json

@@ -27,6 +27,9 @@
     "DataScience.installKernel": "The Jupyter Kernel '{0}' could not be found and needs to be installed in order to execute cells in this notebook.",
     "DataScience.customizeLayout": "Customize Layout",
     "DataScience.warnWhenSelectingKernelWithUnSupportedPythonVersion": "The version of Python associated with the selected kernel is no longer supported. Please consider selecting a different kernel.",
+    "jupyter.kernel.percentPipCondaInstallInsteadOfBang": "Use '%{0} install' instead of '!{0} install'",
+    "jupyter.kernel.pipCondaInstallHoverWarning": "'!{0} install' could install packages into the wrong environment. [More info]({1})",
+    "jupyter.kernel.replacePipCondaInstallCodeAction": "Replace with '%{0} install'",
     "jupyter.kernel.filter.placeholder": "Choose the kernels that are available in the kernel picker.",
     "jupyter.kernel.category.jupyterSession": "Jupyter Session",
     "jupyter.kernel.category.jupyterKernel": "Jupyter Kernel",

src/client/common/utils/localize.ts

@@ -986,6 +986,18 @@ export namespace DataScience {
         'DataScience.fileSeemsToBeInterferingWithKernelStartup',
         "The file '{0}' seems to be overriding built in modules and interfering with the startup of the kernel. Consider renaming the file and starting the kernel again.."
     );
+    export const pipCondaInstallHoverWarning = localize(
+        'jupyter.kernel.pipCondaInstallHoverWarning',
+        "'!{0} install' could install packages into the wrong environment. [More info]({1})"
+    );
+    export const percentPipCondaInstallInsteadOfBang = localize(
+        'jupyter.kernel.percentPipCondaInstallInsteadOfBang',
+        "Use '%{0} install' instead of '!{0} install'"
+    );
+    export const replacePipCondaInstallCodeAction = localize(
+        'jupyter.kernel.replacePipCondaInstallCodeAction',
+        "Replace with '%{0} install'"
+    );
 }
 
 // Skip using vscode-nls and instead just compute our strings based on key values. Key values

src/client/datascience/notebook/diagnosticsProvider.ts

@@ -0,0 +1,234 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { injectable, inject } from 'inversify';
+import {
+    CancellationToken,
+    CodeAction,
+    CodeActionContext,
+    CodeActionKind,
+    CodeActionProvider,
+    DiagnosticSeverity,
+    DiagnosticCollection,
+    languages,
+    NotebookCell,
+    NotebookCellKind,
+    NotebookDocument,
+    Position,
+    Range,
+    Selection,
+    TextDocument,
+    Uri,
+    WorkspaceEdit,
+    Hover,
+    HoverProvider
+} from 'vscode';
+import { IExtensionSyncActivationService } from '../../activation/types';
+import { IDocumentManager, IVSCodeNotebook } from '../../common/application/types';
+import { PYTHON_LANGUAGE } from '../../common/constants';
+import { disposeAllDisposables } from '../../common/helpers';
+import { IDisposable, IDisposableRegistry } from '../../common/types';
+import { DataScience } from '../../common/utils/localize';
+import { JupyterNotebookView } from './constants';
+
+type CellUri = string;
+type CellVersion = number;
+
+const pipMessage = DataScience.percentPipCondaInstallInsteadOfBang().format('pip');
+const condaMessage = DataScience.percentPipCondaInstallInsteadOfBang().format('conda');
+const diagnosticSource = 'Jupyter';
+
+@injectable()
+export class NotebookCellBangInstallDiagnosticsProvider
+    implements IExtensionSyncActivationService, CodeActionProvider, HoverProvider {
+    public readonly problems: DiagnosticCollection;
+    private readonly disposables: IDisposable[] = [];
+    private readonly notebooksProcessed = new WeakMap<NotebookDocument, Map<CellUri, CellVersion>>();
+    private readonly cellsToProcess = new Set<NotebookCell>();
+    constructor(
+        @inject(IVSCodeNotebook) private readonly notebooks: IVSCodeNotebook,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(IDocumentManager) private readonly documents: IDocumentManager
+    ) {
+        this.problems = languages.createDiagnosticCollection(diagnosticSource);
+        this.disposables.push(this.problems);
+        disposables.push(this);
+    }
+    public dispose() {
+        disposeAllDisposables(this.disposables);
+        this.problems.dispose();
+    }
+    public activate(): void {
+        this.disposables.push(languages.registerCodeActionsProvider(PYTHON_LANGUAGE, this));
+        this.disposables.push(languages.registerHoverProvider(PYTHON_LANGUAGE, this));
+        this.documents.onDidChangeTextDocument(
+            (e) => {
+                const notebook = e.document.notebook;
+                if (notebook?.notebookType !== JupyterNotebookView) {
+                    return;
+                }
+                const cell = notebook.getCells().find((c) => c.document === e.document);
+                if (cell) {
+                    this.analyzeNotebookCell(cell);
+                }
+            },
+            this,
+            this.disposables
+        );
+        this.notebooks.onDidCloseNotebookDocument(
+            (e) => {
+                this.problems.delete(e.uri);
+                const cells = this.notebooksProcessed.get(e);
+                this.notebooksProcessed.delete(e);
+                if (!cells) {
+                    return;
+                }
+                cells.forEach((_, uri) => this.problems.delete(Uri.parse(uri)));
+            },
+            this,
+            this.disposables
+        );
+
+        this.notebooks.onDidOpenNotebookDocument((e) => this.analyzeNotebook(e), this, this.disposables);
+        this.notebooks.onDidChangeNotebookDocument(
+            (e) => {
+                if (e.type === 'changeCells') {
+                    const cells = this.notebooksProcessed.get(e.document);
+                    e.changes.forEach((change) => {
+                        change.deletedItems.forEach((cell) => {
+                            cells?.delete(cell.document.uri.toString());
+                        });
+                        change.items.forEach((cell) => this.queueCellForProcessing(cell));
+                    });
+                }
+            },
+            this,
+            this.disposables
+        );
+        this.notebooks.notebookDocuments.map((e) => this.analyzeNotebook(e));
+    }
+    public provideHover(document: TextDocument, position: Position, _token: CancellationToken) {
+        if (document.notebook?.notebookType !== JupyterNotebookView) {
+            return;
+        }
+        const diagnostics = this.problems.get(document.uri);
+        if (!diagnostics) {
+            return;
+        }
+        const diagnostic = diagnostics.find((d) => d.message === pipMessage || d.message === condaMessage);
+        if (!diagnostic || !diagnostic.range.contains(position)) {
+            return;
+        }
+        const installer = diagnostic.message === pipMessage ? 'pip' : 'conda';
+        return new Hover(
+            DataScience.pipCondaInstallHoverWarning().format(installer, 'https://aka.ms/jupyterCellMagicBangInstall'),
+            diagnostic.range
+        );
+    }
+
+    public provideCodeActions(
+        document: TextDocument,
+        _range: Range | Selection,
+        context: CodeActionContext,
+        _token: CancellationToken
+    ): CodeAction[] {
+        if (document.notebook?.notebookType !== JupyterNotebookView) {
+            return [];
+        }
+        const codeActions: CodeAction[] = [];
+        context.diagnostics.forEach((d) => {
+            if (d.message !== pipMessage && d.message !== condaMessage) {
+                return;
+            }
+            const isPip = d.message === pipMessage;
+            const installer = isPip ? 'pip' : 'conda';
+            const codeAction = new CodeAction(
+                DataScience.replacePipCondaInstallCodeAction().format(installer),
+                CodeActionKind.QuickFix
+            );
+            codeAction.isPreferred = true;
+            codeAction.diagnostics = [d];
+            const edit = new WorkspaceEdit();
+            edit.replace(
+                document.uri,
+                new Range(d.range.start, new Position(d.range.start.line, d.range.start.character + 1)),
+                '%'
+            );
+            codeAction.edit = edit;
+            codeActions.push(codeAction);
+        });
+        return codeActions;
+    }
+    private analyzeNotebook(notebook: NotebookDocument): void {
+        if (notebook.notebookType !== JupyterNotebookView) {
+            return;
+        }
+        // Process just the first 100 cells to avoid blocking the UI.
+        notebook.getCells().forEach((cell, i) => (i < 100 ? this.queueCellForProcessing(cell) : undefined));
+    }
+
+    private queueCellForProcessing(cell: NotebookCell): void {
+        this.cellsToProcess.add(cell);
+        this.analyzeNotebookCells();
+    }
+    private analyzeNotebookCells() {
+        if (this.cellsToProcess.size === 0) {
+            return;
+        }
+        const cell = this.cellsToProcess.values().next().value;
+        this.cellsToProcess.delete(cell);
+        this.analyzeNotebookCell(cell);
+        // Schedule processing of next cell (this way we dont chew CPU and block the UI).
+        setTimeout(() => this.analyzeNotebookCells(), 0);
+    }
+    private analyzeNotebookCell(cell: NotebookCell) {
+        if (
+            cell.kind !== NotebookCellKind.Code ||
+            cell.document.languageId !== PYTHON_LANGUAGE ||
+            cell.notebook.isClosed ||
+            cell.document.isClosed
+        ) {
+            return;
+        }
+        // If we've already process this same cell, and the version is the same, then we don't need to do anything.
+        if (this.notebooksProcessed.get(cell.notebook)?.get(cell.document.uri.toString()) === cell.document.version) {
+            return;
+        }
+
+        this.problems.delete(cell.document.uri);
+        const cellsUrisWithProblems = this.notebooksProcessed.get(cell.notebook) || new Map<CellUri, CellVersion>();
+        cellsUrisWithProblems.set(cell.document.uri.toString(), cell.document.version);
+        this.notebooksProcessed.set(cell.notebook, cellsUrisWithProblems);
+
+        // For perf reasons, process just the first 50 lines.
+        for (let i = 0; i < Math.min(cell.document.lineCount, 50); i++) {
+            const line = cell.document.lineAt(i);
+            const text = line.text;
+            if (text.trim().startsWith('!pip install')) {
+                const startPos = text.indexOf('!');
+                const endPos = text.indexOf('l');
+                const range = new Range(line.lineNumber, startPos, line.lineNumber, endPos + 2);
+                this.problems.set(cell.document.uri, [
+                    {
+                        message: pipMessage,
+                        range,
+                        severity: DiagnosticSeverity.Error,
+                        source: diagnosticSource
+                    }
+                ]);
+            } else if (text.trim().startsWith('!conda install')) {
+                const startPos = text.indexOf('!');
+                const endPos = text.indexOf('l');
+                const range = new Range(line.lineNumber, startPos, line.lineNumber, endPos + 2);
+                this.problems.set(cell.document.uri, [
+                    {
+                        message: condaMessage,
+                        range,
+                        severity: DiagnosticSeverity.Error,
+                        source: diagnosticSource
+                    }
+                ]);
+            }
+        }
+    }
+}

src/client/datascience/notebook/serviceRegistry.ts

@@ -25,6 +25,7 @@ import { CellOutputDisplayIdTracker } from '../jupyter/kernels/cellDisplayIdTrac
 import { IntellisenseProvider } from './intellisense/intellisenseProvider';
 import { KernelFilterUI } from './kernelFilter/kernelFilterUI';
 import { KernelFilterService } from './kernelFilter/kernelFilterService';
+import { NotebookCellBangInstallDiagnosticsProvider } from './diagnosticsProvider';
 
 export function registerTypes(serviceManager: IServiceManager) {
     serviceManager.addSingleton<IExtensionSingleActivationService>(
@@ -51,6 +52,10 @@ export function registerTypes(serviceManager: IServiceManager) {
         IExtensionSingleActivationService,
         NotebookCellLanguageService
     );
+    serviceManager.addSingleton<IExtensionSyncActivationService>(
+        IExtensionSyncActivationService,
+        NotebookCellBangInstallDiagnosticsProvider
+    );
     serviceManager.addSingleton<NotebookCellLanguageService>(NotebookCellLanguageService, NotebookCellLanguageService);
     serviceManager.addSingleton<PythonKernelCompletionProvider>(
         PythonKernelCompletionProvider,

src/test/datascience/notebook/diagnosticProvider.vscode.test.ts

@@ -0,0 +1,86 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+/* eslint-disable @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires */
+import { assert } from 'chai';
+import { DataScience } from '../../../client/common/utils/localize';
+import { IVSCodeNotebook } from '../../../client/common/application/types';
+import { traceInfo } from '../../../client/common/logger';
+import { IDisposable } from '../../../client/common/types';
+import { captureScreenShot, IExtensionTestApi, waitForCondition } from '../../common';
+import { initialize } from '../../initialize';
+import {
+    closeNotebooksAndCleanUpAfterTests,
+    insertCodeCell,
+    createEmptyPythonNotebook,
+    workAroundVSCodeNotebookStartPages
+} from './helper';
+import { NotebookCellBangInstallDiagnosticsProvider } from '../../../client/datascience/notebook/diagnosticsProvider';
+import { NotebookDocument, Range } from 'vscode';
+import { IExtensionSyncActivationService } from '../../../client/activation/types';
+
+/* eslint-disable @typescript-eslint/no-explicit-any, no-invalid-this */
+suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
+    let api: IExtensionTestApi;
+    const disposables: IDisposable[] = [];
+    let vscodeNotebook: IVSCodeNotebook;
+    let diagnosticProvider: NotebookCellBangInstallDiagnosticsProvider;
+    let activeNotebook: NotebookDocument;
+    setup(async function () {
+        try {
+            traceInfo(`Start Test ${this.currentTest?.title}`);
+            api = await initialize();
+            await workAroundVSCodeNotebookStartPages();
+            vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
+            diagnosticProvider = api.serviceContainer
+                .getAll<NotebookCellBangInstallDiagnosticsProvider>(IExtensionSyncActivationService)
+                .find((item) => item instanceof NotebookCellBangInstallDiagnosticsProvider)!;
+            await createEmptyPythonNotebook(disposables);
+            activeNotebook = vscodeNotebook.activeNotebookEditor!.document;
+            assert.isOk(activeNotebook, 'No active notebook');
+            traceInfo(`Start Test (completed) ${this.currentTest?.title}`);
+        } catch (e) {
+            await captureScreenShot(this.currentTest?.title || 'unknown');
+            throw e;
+        }
+    });
+    teardown(async function () {
+        traceInfo(`Ended Test ${this.currentTest?.title}`);
+        await closeNotebooksAndCleanUpAfterTests(disposables);
+        traceInfo(`Ended Test (completed) ${this.currentTest?.title}`);
+    });
+    test('Show error for pip install', async () => {
+        await insertCodeCell('!pip install xyz', { index: 0 });
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
+
+        await waitForCondition(
+            async () => (diagnosticProvider.problems.get(cell.document.uri) || []).length > 0,
+            5_000,
+            'No problems detected'
+        );
+        const problem = diagnosticProvider.problems.get(cell.document.uri)![0];
+        assert.equal(problem.message, DataScience.percentPipCondaInstallInsteadOfBang().format('pip'));
+        assert.isTrue(
+            problem.range.isEqual(new Range(0, 0, 0, 12)),
+            `Range is not as expected ${problem.range.toString()}`
+        );
+    });
+    test('Show error for conda install', async () => {
+        await insertCodeCell('!conda install xyz', { index: 0 });
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
+
+        await waitForCondition(
+            async () => (diagnosticProvider.problems.get(cell.document.uri) || []).length > 0,
+            5_000,
+            'No problems detected'
+        );
+        const problem = diagnosticProvider.problems.get(cell.document.uri)![0];
+        assert.equal(problem.message, DataScience.percentPipCondaInstallInsteadOfBang().format('conda'));
+        assert.isTrue(
+            problem.range.isEqual(new Range(0, 0, 0, 14)),
+            `Range is not as expected ${problem.range.toString()}`
+        );
+    });
+});