Support for semantic tokens in notebooks

news/2 Fixes/6799.md

@@ -0,0 +1 @@
+Fix semantic colorization in notebooks.

package.json

@@ -1915,7 +1915,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.15",
         "ansi-to-html": "^0.6.7",
         "arch": "^2.1.0",
         "bootstrap": "^4.3.1",

src/client/datascience/notebook/intellisense/intellisenseProvider.ts

@@ -22,7 +22,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
@@ -114,7 +114,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);
+            }
         }
     }
 
@@ -124,11 +129,11 @@ 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;
     }

src/client/datascience/notebook/intellisense/languageServer.ts

@@ -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);
         }

src/test/common.ts

@@ -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);
     });

src/test/datascience/notebook/intellisense/semanticTokens.vscode.test.ts

@@ -0,0 +1,92 @@
+// 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 () => {
+        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');
+    });
+});