Experimental cell execution analysis

package.json

@@ -775,6 +775,29 @@
                 "icon": "$(vm)",
                 "enablement": "true",
                 "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.gatherCells",
+                "title": "Gather code",
+                "shortTitle": "Gather",
+                "icon": "$(gather)",
+                "enablement": "!jupyter.webExtension",
+                "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.selectSuccessorCells",
+                "title": "Select Successors",
+                "shortTitle": "Select",
+                "icon": "$(arrow-down)",
+                "enablement": "!jupyter.webExtension",
+                "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.debugCellSymbols",
+                "title": "Debug Cell Symbols",
+                "icon": "$(debug-alt-small)",
+                "enablement": "!jupyter.webExtension",
+                "category": "Jupyter"
             }
         ],
         "submenus": [
@@ -925,6 +948,16 @@
                     "command": "jupyter.runByLineStop",
                     "when": "notebookCellResource in jupyter.notebookeditor.runByLineCells && notebookCellToolbarLocation == right",
                     "group": "inline/cell@0"
+                },
+                {
+                    "command": "jupyter.gatherCells",
+                    "when": "notebookType == 'jupyter-notebook' && isWorkspaceTrusted && config.jupyter.executionAnalysis.enabled",
+                    "group": "executionAnalysis@0"
+                },
+                {
+                    "command": "jupyter.selectSuccessorCells",
+                    "when": "notebookType == 'jupyter-notebook' && isWorkspaceTrusted && config.jupyter.executionAnalysis.enabled",
+                    "group": "executionAnalysis@0"
                 }
             ],
             "notebook/cell/execute": [
@@ -1365,6 +1398,18 @@
                 {
                     "command": "jupyter.clearSavedJupyterUris",
                     "title": "%jupyter.command.jupyter.clearSavedJupyterUris.title%"
+                },
+                {
+                    "command": "jupyter.gatherCells",
+                    "when": "config.jupyter.executionAnalysis.enabled"
+                },
+                {
+                    "command": "jupyter.selectSuccessorCells",
+                    "when": "config.jupyter.executionAnalysis.enabled"
+                },
+                {
+                    "command": "jupyter.debugCellSymbols",
+                    "when": "config.jupyter.executionAnalysis.enabled"
                 }
             ],
             "debug/variables/context": [
@@ -1874,6 +1919,12 @@
                     "default": false,
                     "markdownDescription": "%jupyter.configuration.jupyter.enableExtendedKernelCompletions.markdownDescription%",
                     "scope": "application"
+                },
+                "jupyter.executionAnalysis.enabled": {
+                    "type": "boolean",
+                    "default": false,
+                    "description": "Experimental feature to enable execution analysis in notebooks",
+                    "scope": "application"
                 }
             }
         },

src/extension.node.ts

@@ -92,6 +92,10 @@ import { IInterpreterPackages } from './platform/interpreter/types';
 import { homedir, platform, arch, userInfo } from 'os';
 import { getUserHomeDir } from './platform/common/utils/platform.node';
 import { homePath } from './platform/common/platform/fs-paths.node';
+import {
+    activate as activateExecutionAnalysis,
+    deactivate as deactivateExecutionAnalysis
+} from './standalone/executionAnalysis/extension';
 
 durations.codeLoadingTime = stopWatch.elapsedTime;
 
@@ -151,6 +155,8 @@ export function deactivate(): Thenable<void> {
         }
     }
 
+    deactivateExecutionAnalysis();
+
     return Promise.resolve();
 }
 
@@ -187,6 +193,10 @@ async function activateUnsafe(
         startupDurations.endActivateTime = startupStopWatch.elapsedTime;
         activationDeferred.resolve();
 
+        //===============================================
+        // dynamically load standalone plugins
+        activateExecutionAnalysis(context).then(noop, noop);
+
         const api = buildApi(activationPromise, serviceManager, serviceContainer, context);
         return [api, activationPromise, serviceContainer];
     } finally {

src/standalone/executionAnalysis/common.ts

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import * as vscode from 'vscode';
+
+export interface Range {
+    /**
+     * The range's start position
+     */
+    start: Position;
+    /**
+     * The range's end position.
+     */
+    end: Position;
+}
+
+export interface Position {
+    /**
+     * Line position in a document (zero-based).
+     */
+    line: number;
+    /**
+     * Character offset on a line in a document (zero-based). Assuming that the line is
+     * represented as a string, the `character` value represents the gap between the
+     * `character` and `character + 1`.
+     *
+     * If the character value is greater than the line length it defaults back to the
+     * line length.
+     */
+    character: number;
+}
+
+export interface Location {
+    uri: string;
+    range: Range;
+}
+
+export interface LocationWithReferenceKind extends Location {
+    kind?: string;
+}
+
+export function cellIndexesToRanges(indexes: number[]): vscode.NotebookRange[] {
+    indexes.sort((a, b) => a - b);
+    const first = indexes.shift();
+
+    if (first === undefined) {
+        return [];
+    }
+
+    return indexes
+        .reduce(
+            function (ranges, num) {
+                if (num <= ranges[0][1]) {
+                    ranges[0][1] = num + 1;
+                } else {
+                    ranges.unshift([num, num + 1]);
+                }
+                return ranges;
+            },
+            [[first, first + 1]]
+        )
+        .reverse()
+        .map((val) => new vscode.NotebookRange(val[0], val[1]));
+}
+
+export function findNotebook(document: vscode.TextDocument): vscode.NotebookDocument | undefined {
+    return vscode.workspace.notebookDocuments.find(
+        (doc) => doc.uri.authority === document.uri.authority && doc.uri.path === document.uri.path
+    );
+}
+
+// eslint-disable-next-line no-empty,@typescript-eslint/no-empty-function
+export function noop() {}
+
+export const PylanceExtension = 'ms-python.vscode-pylance';

src/standalone/executionAnalysis/extension.ts

@@ -0,0 +1,108 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import * as vscode from 'vscode';
+import { activatePylance } from './pylance';
+import { findNotebook, noop } from './common';
+import { SymbolsTracker } from './symbols';
+
+export async function activate(context: vscode.ExtensionContext): Promise<void> {
+    const optInto = vscode.workspace.getConfiguration('jupyter').get<boolean>('executionAnalysis.enabled');
+    if (!optInto) {
+        return;
+    }
+
+    const referencesProvider = await activatePylance();
+    if (!referencesProvider) {
+        vscode.window
+            .showErrorMessage('Could not get references provider from language server, Pylance prerelease required.')
+            .then(noop, noop);
+        return;
+    }
+
+    const symbolsManager = new SymbolsTracker(referencesProvider);
+    context.subscriptions.push(symbolsManager);
+
+    context.subscriptions.push(
+        vscode.commands.registerCommand(
+            'jupyter.selectSuccessorCells',
+            async (cell: vscode.NotebookCell | undefined) => {
+                const doc =
+                    vscode.workspace.textDocuments.find(
+                        (doc) => doc.uri.toString() === cell?.document.uri.toString()
+                    ) ?? vscode.window.activeTextEditor?.document;
+                if (!doc) {
+                    return;
+                }
+
+                const notebook = findNotebook(doc);
+                if (!notebook) {
+                    return;
+                }
+                const cells = notebook.getCells();
+                const currentCell = cells.find((cell) => cell.document.uri.toString() === doc.uri.toString());
+                if (!currentCell) {
+                    return;
+                }
+
+                await symbolsManager.selectSuccessorCells(notebook, currentCell);
+            }
+        )
+    );
+
+    context.subscriptions.push(
+        vscode.commands.registerCommand('jupyter.gatherCells', async (cell: vscode.NotebookCell | undefined) => {
+            const doc =
+                vscode.workspace.textDocuments.find((doc) => doc.uri.toString() === cell?.document.uri.toString()) ??
+                vscode.window.activeTextEditor?.document;
+            if (!doc) {
+                return;
+            }
+
+            const notebook = findNotebook(doc);
+            if (!notebook) {
+                return;
+            }
+            const cells = notebook.getCells();
+            const currentCell = cells.find((cell) => cell.document.uri.toString() === doc.uri.toString());
+            if (!currentCell) {
+                return;
+            }
+
+            const gatheredCells = (await symbolsManager.gatherCells(notebook, currentCell)) as vscode.NotebookCell[];
+            if (gatheredCells) {
+                // console.log(gatheredCells?.map(cell => `${cell.index}:\n ${cell.document.getText()}\n`));
+
+                const nbCells = gatheredCells.map((cell) => {
+                    return new vscode.NotebookCellData(
+                        vscode.NotebookCellKind.Code,
+                        cell.document.getText(),
+                        cell.document.languageId
+                    );
+                });
+                const doc = await vscode.workspace.openNotebookDocument(
+                    'jupyter-notebook',
+                    new vscode.NotebookData(nbCells)
+                );
+                await vscode.window.showNotebookDocument(doc, {
+                    viewColumn: 1,
+                    preserveFocus: true
+                });
+            }
+        })
+    );
+
+    context.subscriptions.push(
+        vscode.commands.registerCommand('jupyter.debugCellSymbols', async () => {
+            const notebookEditor = vscode.window.activeNotebookEditor;
+            if (notebookEditor) {
+                await symbolsManager.debugSymbols(notebookEditor.notebook);
+            }
+        })
+    );
+}
+
+// This method is called when your extension is deactivated
+export function deactivate() {
+    noop();
+}

src/standalone/executionAnalysis/pylance.ts

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import * as vscode from 'vscode';
+import { BaseLanguageClient } from 'vscode-languageclient';
+import { LocationWithReferenceKind, PylanceExtension, noop } from './common';
+
+export interface ILanguageServerFolder {
+    path: string;
+    version: string; // SemVer, in string form to avoid cross-extension type issues.
+}
+
+export interface INotebookLanguageClient {
+    registerJupyterPythonPathFunction(func: (uri: vscode.Uri) => Promise<string | undefined>): void;
+    registerGetNotebookUriForTextDocumentUriFunction(
+        func: (textDocumentUri: vscode.Uri) => vscode.Uri | undefined
+    ): void;
+    getCompletionItems(
+        document: vscode.TextDocument,
+        position: vscode.Position,
+        context: vscode.CompletionContext,
+        token: vscode.CancellationToken
+    ): Promise<vscode.CompletionItem[] | vscode.CompletionList | undefined>;
+    getReferences(
+        textDocument: vscode.TextDocument,
+        position: vscode.Position,
+        options: {
+            includeDeclaration: boolean;
+        },
+        token: vscode.CancellationToken
+    ): Promise<LocationWithReferenceKind[] | null | undefined>;
+}
+
+export interface LSExtensionApi {
+    languageServerFolder?(): Promise<ILanguageServerFolder>;
+    client?: {
+        isEnabled(): boolean;
+        start(): Promise<void>;
+        stop(): Promise<void>;
+    };
+    notebook?: INotebookLanguageClient;
+}
+
+export interface PythonApi {
+    readonly pylance?: {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        createClient(...args: any[]): BaseLanguageClient;
+        start(client: BaseLanguageClient): Promise<void>;
+        stop(client: BaseLanguageClient): Promise<void>;
+    };
+}
+
+export async function runPylance(pylanceExtension: vscode.Extension<LSExtensionApi>) {
+    const pylanceApi = await pylanceExtension.activate();
+    return pylanceApi;
+}
+
+let _client: INotebookLanguageClient | undefined;
+export async function activatePylance(): Promise<INotebookLanguageClient | undefined> {
+    const pylanceExtension = vscode.extensions.getExtension(PylanceExtension);
+    if (!pylanceExtension) {
+        return undefined;
+    }
+
+    if (_client) {
+        return _client;
+    }
+
+    return new Promise((resolve, reject) => {
+        runPylance(pylanceExtension)
+            .then(async (client) => {
+                if (!client) {
+                    vscode.window.showErrorMessage('Could not start Pylance').then(noop, noop);
+                    reject();
+                    return;
+                }
+
+                if (client.client) {
+                    await client.client.start();
+                }
+
+                _client = client.notebook;
+                resolve(client.notebook);
+            })
+            .then(noop, noop);
+    });
+}