Support for LSP notebooks experiment

package.json

@@ -906,6 +906,15 @@
                     "scope": "machine-overridable",
                     "type": "string"
                 },
+                "python.pylanceLspNotebooksEnabled": {
+                    "type": "boolean",
+                    "default": false,
+                    "description": "Determines if Pylance's experimental LSP notebooks support is used or not.",
+                    "scope": "machine",
+                    "tags": [
+                        "experimental"
+                    ]
+                },
                 "python.sortImports.args": {
                     "default": [],
                     "description": "Arguments passed in. Each argument is a separate item in the array.",

src/client/activation/jedi/languageClientMiddleware.ts

@@ -0,0 +1,13 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { IServiceContainer } from '../../ioc/types';
+import { LanguageClientMiddleware } from '../languageClientMiddleware';
+import { LanguageServerType } from '../types';
+
+export class JediLanguageClientMiddleware extends LanguageClientMiddleware {
+    public constructor(serviceContainer: IServiceContainer, serverVersion?: string) {
+        super(serviceContainer, LanguageServerType.Jedi, serverVersion);
+        this.setupHidingMiddleware(serviceContainer);
+    }
+}

src/client/activation/jedi/languageServerProxy.ts

@@ -14,7 +14,7 @@ import { IInterpreterPathService, Resource } from '../../common/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
-import { LanguageClientMiddleware } from '../languageClientMiddleware';
+import { JediLanguageClientMiddleware } from './languageClientMiddleware';
 import { ProgressReporting } from '../progress';
 import { ILanguageClientFactory, ILanguageServerProxy } from '../types';
 import { killPid } from '../../common/process/rawProcessApis';
@@ -57,7 +57,8 @@ export class JediLanguageServerProxy implements ILanguageServerProxy {
         options: LanguageClientOptions,
     ): Promise<void> {
         this.lsVersion =
-            (options.middleware ? (<LanguageClientMiddleware>options.middleware).serverVersion : undefined) ?? '0.19.3';
+            (options.middleware ? (<JediLanguageClientMiddleware>options.middleware).serverVersion : undefined) ??
+            '0.19.3';
 
         this.languageClient = await this.factory.createLanguageClient(resource, interpreter, options);
         this.registerHandlers();

src/client/activation/jedi/manager.ts

@@ -14,21 +14,16 @@ import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
 import { Commands } from '../commands';
-import { LanguageClientMiddleware } from '../languageClientMiddleware';
-import {
-    ILanguageServerAnalysisOptions,
-    ILanguageServerManager,
-    ILanguageServerProxy,
-    LanguageServerType,
-} from '../types';
+import { JediLanguageClientMiddleware } from './languageClientMiddleware';
+import { ILanguageServerAnalysisOptions, ILanguageServerManager, ILanguageServerProxy } from '../types';
 import { traceDecoratorError, traceDecoratorVerbose, traceVerbose } from '../../logging';
 
 export class JediLanguageServerManager implements ILanguageServerManager {
     private resource!: Resource;
 
     private interpreter: PythonEnvironment | undefined;
 
-    private middleware: LanguageClientMiddleware | undefined;
+    private middleware: JediLanguageClientMiddleware | undefined;
 
     private disposables: IDisposable[] = [];
 
@@ -132,7 +127,7 @@ export class JediLanguageServerManager implements ILanguageServerManager {
     @traceDecoratorVerbose('Starting language server')
     protected async startLanguageServer(): Promise<void> {
         const options = await this.analysisOptions.getAnalysisOptions();
-        this.middleware = new LanguageClientMiddleware(this.serviceContainer, LanguageServerType.Jedi, this.lsVersion);
+        this.middleware = new JediLanguageClientMiddleware(this.serviceContainer, this.lsVersion);
         options.middleware = this.middleware;
 
         // Make sure the middleware is connected if we restart and we we're already connected.

src/client/activation/languageClientMiddleware.ts

@@ -14,31 +14,45 @@ import { createHidingMiddleware } from '@vscode/jupyter-lsp-middleware';
 export class LanguageClientMiddleware extends LanguageClientMiddlewareBase {
     public constructor(serviceContainer: IServiceContainer, serverType: LanguageServerType, serverVersion?: string) {
         super(serviceContainer, serverType, sendTelemetryEvent, serverVersion);
+    }
 
-        if (serverType === LanguageServerType.None) {
-            return;
-        }
-
+    /**
+     * Creates the HidingMiddleware if needed and sets up code to do so if needed after
+     * Jupyter is installed.
+     *
+     * This method should be called from the constructor of derived classes. It is separated
+     * from the constructor to allow derived classes to initialize before it is called.
+     */
+    protected setupHidingMiddleware(serviceContainer: IServiceContainer) {
         const jupyterDependencyManager = serviceContainer.get<IJupyterExtensionDependencyManager>(
             IJupyterExtensionDependencyManager,
         );
         const disposables = serviceContainer.get<IDisposableRegistry>(IDisposableRegistry) || [];
         const extensions = serviceContainer.get<IExtensions>(IExtensions);
 
         // Enable notebook support if jupyter support is installed
-        if (jupyterDependencyManager && jupyterDependencyManager.isJupyterExtensionInstalled) {
+        if (this.shouldCreateHidingMiddleware(jupyterDependencyManager)) {
             this.notebookAddon = createHidingMiddleware();
         }
+
         disposables.push(
-            extensions?.onDidChange(() => {
-                if (jupyterDependencyManager) {
-                    if (this.notebookAddon && !jupyterDependencyManager.isJupyterExtensionInstalled) {
-                        this.notebookAddon = undefined;
-                    } else if (!this.notebookAddon && jupyterDependencyManager.isJupyterExtensionInstalled) {
-                        this.notebookAddon = createHidingMiddleware();
-                    }
-                }
+            extensions?.onDidChange(async () => {
+                await this.onExtensionChange(jupyterDependencyManager);
             }),
         );
     }
+
+    protected shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
+        return jupyterDependencyManager && jupyterDependencyManager.isJupyterExtensionInstalled;
+    }
+
+    protected async onExtensionChange(jupyterDependencyManager: IJupyterExtensionDependencyManager): Promise<void> {
+        if (jupyterDependencyManager) {
+            if (this.notebookAddon && !this.shouldCreateHidingMiddleware(jupyterDependencyManager)) {
+                this.notebookAddon = undefined;
+            } else if (!this.notebookAddon && this.shouldCreateHidingMiddleware(jupyterDependencyManager)) {
+                this.notebookAddon = createHidingMiddleware();
+            }
+        }
+    }
 }

src/client/activation/languageClientMiddlewareBase.ts

@@ -86,7 +86,9 @@ export class LanguageClientMiddlewareBase implements Middleware {
                     const settingDict: LSPObject & { pythonPath: string; _envPYTHONPATH: string } = settings[
                         i
                     ] as LSPObject & { pythonPath: string; _envPYTHONPATH: string };
-                    settingDict.pythonPath = configService.getSettings(uri).pythonPath;
+
+                    settingDict.pythonPath =
+                        (await this.getPythonPathOverride(uri)) ?? configService.getSettings(uri).pythonPath;
 
                     const env = await envService.getEnvironmentVariables(uri);
                     const envPYTHONPATH = env.PYTHONPATH;
@@ -100,6 +102,11 @@ export class LanguageClientMiddlewareBase implements Middleware {
         },
     };
 
+    // eslint-disable-next-line class-methods-use-this
+    protected async getPythonPathOverride(_uri: Uri | undefined): Promise<string | undefined> {
+        return undefined;
+    }
+
     private get connected(): Promise<boolean> {
         return this.connectedPromise.promise;
     }

src/client/activation/node/analysisOptions.ts

@@ -6,10 +6,15 @@ import { IWorkspaceService } from '../../common/application/types';
 
 import { LanguageServerAnalysisOptionsBase } from '../common/analysisOptions';
 import { ILanguageServerOutputChannel } from '../types';
+import { LspNotebooksExperiment } from './lspNotebooksExperiment';
 
 export class NodeLanguageServerAnalysisOptions extends LanguageServerAnalysisOptionsBase {
     // eslint-disable-next-line @typescript-eslint/no-useless-constructor
-    constructor(lsOutputChannel: ILanguageServerOutputChannel, workspace: IWorkspaceService) {
+    constructor(
+        lsOutputChannel: ILanguageServerOutputChannel,
+        workspace: IWorkspaceService,
+        private readonly lspNotebooksExperiment: LspNotebooksExperiment,
+    ) {
         super(lsOutputChannel, workspace);
     }
 
@@ -18,6 +23,7 @@ export class NodeLanguageServerAnalysisOptions extends LanguageServerAnalysisOpt
         return ({
             experimentationSupport: true,
             trustedWorkspaceSupport: true,
+            lspNotebooksSupport: this.lspNotebooksExperiment.isInNotebooksExperiment(),
         } as unknown) as LanguageClientOptions;
     }
 }

src/client/activation/node/languageClientMiddleware.ts

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { Uri } from 'vscode';
+import { IJupyterExtensionDependencyManager } from '../../common/application/types';
+import { IServiceContainer } from '../../ioc/types';
+import { JupyterExtensionIntegration } from '../../jupyter/jupyterIntegration';
+import { traceLog } from '../../logging';
+import { LanguageClientMiddleware } from '../languageClientMiddleware';
+
+import { LanguageServerType } from '../types';
+
+import { LspNotebooksExperiment } from './lspNotebooksExperiment';
+
+export class NodeLanguageClientMiddleware extends LanguageClientMiddleware {
+    private readonly lspNotebooksExperiment: LspNotebooksExperiment;
+
+    public constructor(serviceContainer: IServiceContainer, serverVersion?: string) {
+        super(serviceContainer, LanguageServerType.Node, serverVersion);
+
+        this.lspNotebooksExperiment = serviceContainer.get<LspNotebooksExperiment>(LspNotebooksExperiment);
+        this.setupHidingMiddleware(serviceContainer);
+    }
+
+    protected shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
+        return (
+            super.shouldCreateHidingMiddleware(jupyterDependencyManager) &&
+            !this.lspNotebooksExperiment.isInNotebooksExperiment()
+        );
+    }
+
+    protected async onExtensionChange(jupyterDependencyManager: IJupyterExtensionDependencyManager): Promise<void> {
+        if (jupyterDependencyManager && jupyterDependencyManager.isJupyterExtensionInstalled) {
+            await this.lspNotebooksExperiment.onJupyterInstalled();
+        }
+
+        super.onExtensionChange(jupyterDependencyManager);
+    }
+
+    protected async getPythonPathOverride(uri: Uri | undefined): Promise<string | undefined> {
+        if (!uri || !this.lspNotebooksExperiment.isInNotebooksExperiment()) {
+            return undefined;
+        }
+
+        const jupyterExtensionIntegration = this.serviceContainer?.get<JupyterExtensionIntegration>(
+            JupyterExtensionIntegration,
+        );
+        const jupyterPythonPathFunction = jupyterExtensionIntegration?.getJupyterPythonPathFunction();
+        if (!jupyterPythonPathFunction) {
+            return undefined;
+        }
+
+        const result = await jupyterPythonPathFunction(uri);
+
+        if (result) {
+            traceLog(`Jupyter provided interpreter path override: ${result}`);
+        }
+
+        return result;
+    }
+}