Skip to content

Commit

Permalink
Experimental cell execution analysis. (#14507)
Browse files Browse the repository at this point in the history
  • Loading branch information
rebornix authored Oct 13, 2023
1 parent 10fb995 commit b5a9010
Show file tree
Hide file tree
Showing 6 changed files with 806 additions and 0 deletions.
51 changes: 51 additions & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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"
}
}
},
Expand Down
10 changes: 10 additions & 0 deletions src/extension.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -151,6 +155,8 @@ export function deactivate(): Thenable<void> {
}
}

deactivateExecutionAnalysis();

return Promise.resolve();
}

Expand Down Expand Up @@ -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 {
Expand Down
75 changes: 75 additions & 0 deletions src/standalone/executionAnalysis/common.ts
Original file line number Diff line number Diff line change
@@ -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';
108 changes: 108 additions & 0 deletions src/standalone/executionAnalysis/extension.ts
Original file line number Diff line number Diff line change
@@ -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();
}
87 changes: 87 additions & 0 deletions src/standalone/executionAnalysis/pylance.ts
Original file line number Diff line number Diff line change
@@ -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);
});
}
Loading

0 comments on commit b5a9010

Please sign in to comment.