Skip to content

Commit

Permalink
Support for semantic tokens in notebooks (#8001)
Browse files Browse the repository at this point in the history
* Interpreter id was getting dupes

* Add a test update waitForCondition to not continually spam for the promise

* JediLSP isn't a thing

* Update to latest npm middleware

* Add news entry

* Update to latest and support intellisense in interactive

* Disable semantic token test
  • Loading branch information
rchiodo authored Oct 25, 2021
1 parent 4da1d9e commit 9c830c5
Show file tree
Hide file tree
Showing 9 changed files with 157 additions and 21 deletions.
1 change: 1 addition & 0 deletions news/2 Fixes/6799.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix semantic colorization in notebooks.
6 changes: 3 additions & 3 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -1962,7 +1962,7 @@
"@jupyterlab/services": "^6.1.17",
"@lumino/widgets": "^1.28.0",
"@nteract/messaging": "^7.0.0",
"@vscode/jupyter-lsp-middleware": "^0.2.14",
"@vscode/jupyter-lsp-middleware": "^0.2.16",
"ansi-to-html": "^0.6.7",
"arch": "^2.1.0",
"bootstrap": "^4.3.1",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
private _notebookDocument: NotebookDocument | undefined;
private executionPromise: Promise<boolean> | undefined;
private _notebookEditor: NotebookEditor | undefined;
private _inputUri: Uri | undefined;

constructor(
private readonly applicationShell: IApplicationShell,
Expand Down Expand Up @@ -195,7 +196,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
const controllerId = preferredController ? `${JVSC_EXTENSION_ID}/${preferredController.id}` : undefined;
traceInfo(`Starting interactive window with controller ID ${controllerId}`);
const hasOwningFile = this.owner !== undefined;
const { notebookEditor } = ((await this.commandManager.executeCommand(
const { inputUri, notebookEditor } = ((await this.commandManager.executeCommand(
'interactive.open',
// Keep focus on the owning file if there is one
{ viewColumn: ViewColumn.Beside, preserveFocus: hasOwningFile },
Expand All @@ -210,6 +211,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
}
this._notebookEditor = notebookEditor;
this._notebookDocument = notebookEditor.document;
this._inputUri = inputUri;
this.internalDisposables.push(
window.onDidChangeActiveNotebookEditor((e) => {
if (e === this._notebookEditor) {
Expand Down Expand Up @@ -335,6 +337,10 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
);
}

public get inputUri() {
return this._inputUri;
}

public dispose() {
this.internalDisposables.forEach((d) => d.dispose());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { IDisposableRegistry } from '../../../common/types';
import { IInterpreterService } from '../../../interpreter/contracts';
import { PythonEnvironment } from '../../../pythonEnvironments/info';
import { getInterpreterId } from '../../../pythonEnvironments/info/interpreter';
import { IInteractiveWindowProvider } from '../../types';
import { isJupyterNotebook } from '../helpers/helpers';
import { INotebookControllerManager } from '../types';
import { VSCodeNotebookController } from '../vscodeNotebookController';
Expand All @@ -22,7 +23,7 @@ import { LanguageServer } from './languageServer';
export class IntellisenseProvider implements IExtensionSyncActivationService {
private servers = new Map<string, Promise<LanguageServer | undefined>>();
private activeInterpreterCache = new Map<string, PythonEnvironment | undefined>();
private interpreterIdCache: Map<PythonEnvironment, string> = new Map<PythonEnvironment, string>();
private interpreterIdCache: Map<string, string> = new Map<string, string>();
private knownControllers: WeakMap<NotebookDocument, VSCodeNotebookController> = new WeakMap<
NotebookDocument,
VSCodeNotebookController
Expand All @@ -33,7 +34,8 @@ export class IntellisenseProvider implements IExtensionSyncActivationService {
@inject(INotebookControllerManager) private readonly notebookControllerManager: INotebookControllerManager,
@inject(IVSCodeNotebook) private readonly notebooks: IVSCodeNotebook,
@inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
@inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
@inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
@inject(IInteractiveWindowProvider) private readonly interactiveWindowProvider: IInteractiveWindowProvider
) {}
public activate() {
// Sign up for kernel change events on notebooks
Expand Down Expand Up @@ -114,7 +116,12 @@ export class IntellisenseProvider implements IExtensionSyncActivationService {
// If the controller is empty, default to the active interpreter
const interpreter =
controller?.connection.interpreter || (await this.interpreterService.getActiveInterpreter(n.uri));
return this.ensureLanguageServer(interpreter, n);
const server = await this.ensureLanguageServer(interpreter, n);

// If we created one, make sure the server thinks this file is open
if (server) {
server.startWatching(n);
}
}
}

Expand All @@ -124,19 +131,31 @@ export class IntellisenseProvider implements IExtensionSyncActivationService {
}

private getInterpreterIdFromCache(interpreter: PythonEnvironment) {
let id = this.interpreterIdCache.get(interpreter);
let id = this.interpreterIdCache.get(interpreter.path);
if (!id) {
// Making an assumption that the id for an interpreter never changes.
id = getInterpreterId(interpreter);
this.interpreterIdCache.set(interpreter, id);
this.interpreterIdCache.set(interpreter.path, id);
}
return id;
}

private getNotebook(uri: Uri): NotebookDocument | undefined {
let notebook = this.notebooks.notebookDocuments.find((n) => arePathsSame(n.uri.fsPath, uri.fsPath));
if (!notebook) {
// Might be an interactive window input
const interactiveWindow = this.interactiveWindowProvider.windows.find(
(w) => w.inputUri?.toString() === uri.toString()
);
notebook = interactiveWindow?.notebookDocument;
}
return notebook;
}

private shouldAllowIntellisense(uri: Uri, interpreterId: string, _interpreterPath: string) {
// We should allow intellisense for a URI when the interpreter matches
// the controller for the uri
const notebook = this.notebooks.notebookDocuments.find((n) => arePathsSame(n.uri.fsPath, uri.fsPath));
const notebook = this.getNotebook(uri);
const controller = notebook
? this.notebookControllerManager.getSelectedNotebookController(notebook)
: undefined;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ export class LanguageServer implements Disposable {
cancellationStrategy: FileBasedCancellationStrategy
): Promise<ServerOptions | undefined> {
const pythonConfig = workspace.getConfiguration('python');
if (pythonConfig && pythonConfig.get<string>('languageServer') === 'JediLSP') {
if (pythonConfig && pythonConfig.get<string>('languageServer') === 'Jedi') {
// Use jedi to start our language server.
return LanguageServer.createJediLSPServerOptions(interpreter);
}
Expand Down
1 change: 1 addition & 0 deletions src/client/datascience/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,7 @@ export interface IInteractiveWindow extends IInteractiveBase {
readonly owner: Resource;
readonly submitters: Uri[];
readonly notebookUri?: Uri;
readonly inputUri?: Uri;
readonly notebookDocument?: NotebookDocument;
readonly readyPromise: Promise<void>;
closed: Event<IInteractiveWindow>;
Expand Down
33 changes: 24 additions & 9 deletions src/test/common.ts
Original file line number Diff line number Diff line change
Expand Up @@ -515,25 +515,40 @@ export async function waitForCondition(
condition: () => Promise<boolean>,
timeoutMs: number,
errorMessage: string | (() => string),
intervalTimeoutMs: number = 10
intervalTimeoutMs: number = 10,
throwOnError: boolean = false
): Promise<void> {
return new Promise<void>(async (resolve, reject) => {
const timeout = setTimeout(() => {
clearTimeout(timeout);
// eslint-disable-next-line @typescript-eslint/no-use-before-define
clearInterval(timer);
clearTimeout(timer);
errorMessage = typeof errorMessage === 'string' ? errorMessage : errorMessage();
console.log(`Test failing --- ${errorMessage}`);
reject(new Error(errorMessage));
}, timeoutMs);
const timer = setInterval(async () => {
if (!(await condition().catch(() => false))) {
return;
let timer: NodeJS.Timer;
const timerFunc = async () => {
let success = false;
try {
success = await condition();
} catch (exc) {
if (throwOnError) {
reject(exc);
}
}
clearTimeout(timeout);
clearInterval(timer);
resolve();
}, intervalTimeoutMs);
if (!success) {
// Start up a timer again, but don't do it until after
// the condition is false.
timer = setTimeout(timerFunc, intervalTimeoutMs);
} else {
clearTimeout(timer);
clearTimeout(timeout);
resolve();
}
};
timer = setTimeout(timerFunc, intervalTimeoutMs);

pendingTimers.push(timer);
pendingTimers.push(timeout);
});
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

/* eslint-disable @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires */
import { assert } from 'chai';
import * as sinon from 'sinon';
import { commands } from 'vscode';
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 { IS_REMOTE_NATIVE_TEST } from '../../../constants';
import { initialize } from '../../../initialize';
import {
canRunNotebookTests,
closeNotebooksAndCleanUpAfterTests,
insertCodeCell,
startJupyterServer,
prewarmNotebooks,
createEmptyPythonNotebook,
defaultNotebookTestTimeout
} from '../helper';

/* eslint-disable @typescript-eslint/no-explicit-any, no-invalid-this */
suite('DataScience - VSCode semantic token tests', function () {
let api: IExtensionTestApi;
const disposables: IDisposable[] = [];
let vscodeNotebook: IVSCodeNotebook;
this.timeout(120_000);
suiteSetup(async function () {
traceInfo(`Start Suite semantic token tests`);
this.timeout(120_000);
api = await initialize();
if (IS_REMOTE_NATIVE_TEST) {
// https://github.com/microsoft/vscode-jupyter/issues/6331
return this.skip();
}
if (!(await canRunNotebookTests())) {
return this.skip();
}
await startJupyterServer();
await prewarmNotebooks();
sinon.restore();
vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
traceInfo(`Start Suite (Completed) semantic token tests`);
});
// Use same notebook without starting kernel in every single test (use one for whole suite).
setup(async function () {
traceInfo(`Start Test ${this.currentTest?.title}`);
sinon.restore();
await startJupyterServer();
await createEmptyPythonNotebook(disposables);
traceInfo(`Start Test (completed) ${this.currentTest?.title}`);
});
teardown(async function () {
traceInfo(`Ended Test ${this.currentTest?.title}`);
if (this.currentTest?.isFailed()) {
await captureScreenShot(this.currentTest?.title);
}
await closeNotebooksAndCleanUpAfterTests(disposables);
traceInfo(`Ended Test (completed) ${this.currentTest?.title}`);
});
suiteTeardown(() => closeNotebooksAndCleanUpAfterTests(disposables));
test('Open a notebook and add a bunch of cells', async function () {
// Skip for now. Need to wait for changes to VS code
this.skip();
await insertCodeCell('import sys\nprint(sys.executable)\na = 1');
await insertCodeCell('\ndef test():\n print("test")\ntest()');
const cell1 = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
const cell2 = vscodeNotebook.activeNotebookEditor?.document.cellAt(1)!;

// Wait for tokens on the first cell (it works with just plain pylance)
await waitForCondition(
async () => {
const promise = commands.executeCommand('vscode.provideDocumentSemanticTokens', cell1.document.uri);
const result = (await promise) as any;
return result && result.data.length > 0;
},
defaultNotebookTestTimeout,
`Tokens never appear for first cell`,
100,
true
);

// Then get tokens on second cell. They should start on line 1. If not this
// means there's a bug
const tokens = (await commands.executeCommand(
'vscode.provideDocumentSemanticTokens',
cell2.document.uri
)) as any;
assert.ok(tokens, 'No tokens found on second cell');
assert.equal(tokens.data[0], 1, 'Tokens not correctly offset');
});
});

0 comments on commit 9c830c5

Please sign in to comment.