Improved api to get Python env associated with Jupyter Notebook (#16332)

* API for Python ext to get Python Env for a Notebook

* Misc changes

* Fixes

* Fix formatting

* Fix formatting

* Fix formatting

src/api.proposed.notebookEnvironment.ts

@@ -0,0 +1,37 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import type { Event, Uri } from 'vscode';
+
+declare module './api' {
+    /**
+     * These types are not required for any other extension, except for the Python extension.
+     * Hence the reason to keep this separate. This way we can keep the API stable for other extensions (which would be the majority case).
+     */
+    export interface Jupyter {
+        /**
+         * This event is triggered when the environment associated with a Jupyter Notebook or Interactive Window changes.
+         * The Uri in the event is the Uri of the Notebook/IW.
+         */
+        onDidChangePythonEnvironment: Event<Uri>;
+        /**
+         * Returns the EnvironmentPath to the Python environment associated with a Jupyter Notebook or Interactive Window.
+         * If the Uri is not associated with a Jupyter Notebook or Interactive Window, then this method returns undefined.
+         * @param uri
+         */
+        getPythonEnvironment(uri: Uri):
+            | undefined
+            | {
+                  /**
+                   * The ID of the environment.
+                   */
+                  readonly id: string;
+                  /**
+                   * Path to environment folder or path to python executable that uniquely identifies an environment. Environments
+                   * lacking a python executable are identified by environment folder paths, whereas other envs can be identified
+                   * using python executable path.
+                   */
+                  readonly path: string;
+              };
+    }
+}

src/extension.node.ts

@@ -127,7 +127,10 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             kernels: {
                 getKernel: () => Promise.resolve(undefined),
                 onDidStart: () => ({ dispose: noop })
-            }
+            },
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            onDidChangePythonEnvironment: undefined as any,
+            getPythonEnvironment: () => undefined
         };
     }
 }

src/extension.web.ts

@@ -117,10 +117,10 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             createJupyterServerCollection: () => {
                 throw new Error('Not Implemented');
             },
-            kernels: {
-                getKernel: () => Promise.resolve(undefined),
-                onDidStart: () => ({ dispose: noop })
-            }
+            kernels: { getKernel: () => Promise.resolve(undefined), onDidStart: () => ({ dispose: noop }) },
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            onDidChangePythonEnvironment: undefined as any,
+            getPythonEnvironment: () => undefined
         };
     }
 }

src/notebooks/notebookEnvironmentService.node.ts

@@ -0,0 +1,188 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { EventEmitter, NotebookDocument, Uri } from 'vscode';
+import * as fs from 'fs-extra';
+import { IControllerRegistration, type IVSCodeNotebookController } from './controllers/types';
+import { IKernelProvider, isRemoteConnection, type IKernel } from '../kernels/types';
+import { DisposableBase } from '../platform/common/utils/lifecycle';
+import { isPythonKernelConnection } from '../kernels/helpers';
+import { logger } from '../platform/logging';
+import { getDisplayPath } from '../platform/common/platform/fs-paths.node';
+import { noop } from '../platform/common/utils/misc';
+import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './types';
+import { getCachedEnvironment, getInterpreterInfo } from '../platform/interpreter/helpers';
+import type { Environment } from '@vscode/python-extension';
+import type { PythonEnvironment } from '../platform/pythonEnvironments/info';
+
+@injectable()
+export class NotebookPythonEnvironmentService extends DisposableBase implements INotebookPythonEnvironmentService {
+    private readonly _onDidChangeEnvironment = this._register(new EventEmitter<Uri>());
+    public readonly onDidChangeEnvironment = this._onDidChangeEnvironment.event;
+
+    private readonly notebookWithRemoteKernelsToMonitor = new WeakSet<NotebookDocument>();
+    private readonly notebookPythonEnvironments = new WeakMap<NotebookDocument, Environment | undefined>();
+    constructor(
+        @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration,
+        @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
+        @inject(INotebookEditorProvider) private readonly notebookEditorProvider: INotebookEditorProvider
+    ) {
+        super();
+        this.monitorRemoteKernelStart();
+        this._register(
+            this.controllerRegistration.onControllerSelected((e) => {
+                if (!isPythonKernelConnection(e.controller.connection)) {
+                    this.notebookWithRemoteKernelsToMonitor.delete(e.notebook);
+                    if (this.notebookPythonEnvironments.has(e.notebook)) {
+                        this.notebookPythonEnvironments.delete(e.notebook);
+                        this._onDidChangeEnvironment.fire(e.notebook.uri);
+                    }
+                    return;
+                }
+
+                if (isRemoteConnection(e.controller.connection)) {
+                    this.notebookWithRemoteKernelsToMonitor.add(e.notebook);
+                } else {
+                    this.notebookWithRemoteKernelsToMonitor.delete(e.notebook);
+                    this.notifyLocalPythonEnvironment(e.notebook, e.controller);
+                }
+            })
+        );
+    }
+
+    public getPythonEnvironment(uri: Uri): Environment | undefined {
+        const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(uri);
+        return notebook ? this.notebookPythonEnvironments.get(notebook) : undefined;
+    }
+
+    private monitorRemoteKernelStart() {
+        const trackKernel = async (e: IKernel) => {
+            if (
+                !this.notebookWithRemoteKernelsToMonitor.has(e.notebook) ||
+                !isRemoteConnection(e.kernelConnectionMetadata) ||
+                !isPythonKernelConnection(e.kernelConnectionMetadata)
+            ) {
+                return;
+            }
+
+            try {
+                const env = await this.resolveRemotePythonEnvironment(e.notebook);
+                if (this.controllerRegistration.getSelected(e.notebook)?.controller !== e.controller) {
+                    logger.trace(
+                        `Remote Python Env for ${getDisplayPath(e.notebook.uri)} not determined as controller changed`
+                    );
+                    return;
+                }
+
+                if (!env) {
+                    logger.trace(
+                        `Remote Python Env for ${getDisplayPath(e.notebook.uri)} not determined as exe is empty`
+                    );
+                    return;
+                }
+
+                this.notebookPythonEnvironments.set(e.notebook, env);
+                this._onDidChangeEnvironment.fire(e.notebook.uri);
+            } catch (ex) {
+                logger.error(`Failed to get Remote Python Env for ${getDisplayPath(e.notebook.uri)}`, ex);
+            }
+        };
+        this._register(this.kernelProvider.onDidCreateKernel(trackKernel));
+        this._register(this.kernelProvider.onDidStartKernel(trackKernel));
+    }
+
+    private notifyLocalPythonEnvironment(notebook: NotebookDocument, controller: IVSCodeNotebookController) {
+        // Empty string is special, means do not use any interpreter at all.
+        // Could be a server started for local machine, github codespaces, azml, 3rd party api, etc
+        const connection = this.kernelProvider.get(notebook)?.kernelConnectionMetadata || controller.connection;
+        const interpreter = connection.interpreter;
+        if (!isPythonKernelConnection(connection) || isRemoteConnection(connection) || !interpreter) {
+            return;
+        }
+
+        const env = getCachedEnvironment(interpreter);
+        if (env) {
+            this.notebookPythonEnvironments.set(notebook, env);
+            this._onDidChangeEnvironment.fire(notebook.uri);
+            return;
+        }
+
+        void this.resolveAndNotifyLocalPythonEnvironment(notebook, controller, interpreter);
+    }
+
+    private async resolveAndNotifyLocalPythonEnvironment(
+        notebook: NotebookDocument,
+        controller: IVSCodeNotebookController,
+        interpreter: PythonEnvironment | Readonly<PythonEnvironment>
+    ) {
+        const env = await getInterpreterInfo(interpreter);
+
+        if (!env) {
+            logger.error(
+                `Failed to get interpreter information for ${getDisplayPath(notebook.uri)} && ${getDisplayPath(
+                    interpreter.uri
+                )}`
+            );
+            return;
+        }
+
+        if (this.controllerRegistration.getSelected(notebook)?.controller !== controller.controller) {
+            logger.trace(`Python Env for ${getDisplayPath(notebook.uri)} not determined as controller changed`);
+            return;
+        }
+
+        this.notebookPythonEnvironments.set(notebook, env);
+        this._onDidChangeEnvironment.fire(notebook.uri);
+    }
+
+    private async resolveRemotePythonEnvironment(notebook: NotebookDocument): Promise<Environment | undefined> {
+        // Empty string is special, means do not use any interpreter at all.
+        // Could be a server started for local machine, github codespaces, azml, 3rd party api, etc
+        const kernel = this.kernelProvider.get(notebook);
+        if (!kernel) {
+            return;
+        }
+        if (!kernel.startedAtLeastOnce) {
+            return;
+        }
+        const execution = this.kernelProvider.getKernelExecution(kernel);
+        const code = `
+import os as _VSCODE_os
+import sys as _VSCODE_sys
+import builtins as _VSCODE_builtins
+
+if _VSCODE_os.path.exists("${__filename}"):
+    _VSCODE_builtins.print(f"EXECUTABLE{_VSCODE_sys.executable}EXECUTABLE")
+
+del _VSCODE_os, _VSCODE_sys, _VSCODE_builtins
+`;
+        const outputs = (await execution.executeHidden(code).catch(noop)) || [];
+        const output = outputs.find((item) => item.output_type === 'stream' && item.name === 'stdout');
+        if (!output || !(output.text || '').toString().includes('EXECUTABLE')) {
+            return;
+        }
+        let text = (output.text || '').toString();
+        text = text.substring(text.indexOf('EXECUTABLE'));
+        const items = text.split('EXECUTABLE').filter((x) => x.trim().length);
+        const executable = items.length ? items[0].trim() : '';
+        if (!executable || !(await fs.pathExists(executable))) {
+            return;
+        }
+        logger.debug(
+            `Remote Interpreter for Notebook URI "${getDisplayPath(notebook.uri)}" is ${getDisplayPath(executable)}`
+        );
+
+        const env = getCachedEnvironment(executable) || (await getInterpreterInfo({ id: executable }));
+
+        if (env) {
+            return env;
+        } else {
+            logger.error(
+                `Failed to get remote interpreter information for ${getDisplayPath(notebook.uri)} && ${getDisplayPath(
+                    executable
+                )}`
+            );
+        }
+    }
+}

src/notebooks/notebookEnvironmentService.web.ts

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { injectable } from 'inversify';
+import { EventEmitter, Uri } from 'vscode';
+import { DisposableBase } from '../platform/common/utils/lifecycle';
+import type { INotebookPythonEnvironmentService } from './types';
+import type { Environment } from '@vscode/python-extension';
+
+@injectable()
+export class NotebookPythonEnvironmentService extends DisposableBase implements INotebookPythonEnvironmentService {
+    private readonly _onDidChangeEnvironment = this._register(new EventEmitter<Uri>());
+    public readonly onDidChangeEnvironment = this._onDidChangeEnvironment.event;
+
+    public getPythonEnvironment(_: Uri): Environment | undefined {
+        return undefined;
+    }
+}

src/notebooks/serviceRegistry.node.ts

@@ -36,10 +36,11 @@ import { NotebookCellLanguageService } from './languages/cellLanguageService';
 import { EmptyNotebookCellLanguageService } from './languages/emptyNotebookCellLanguageService';
 import { NotebookCommandListener } from './notebookCommandListener';
 import { NotebookEditorProvider } from './notebookEditorProvider';
+import { NotebookPythonEnvironmentService } from './notebookEnvironmentService.node';
 import { CellOutputMimeTypeTracker } from './outputs/jupyterCellOutputMimeTypeTracker';
 import { NotebookTracebackFormatter } from './outputs/tracebackFormatter';
 import { InterpreterPackageTracker } from './telemetry/interpreterPackageTracker.node';
-import { INotebookEditorProvider } from './types';
+import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './types';
 
 export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) {
     registerControllerTypes(serviceManager, isDevMode);
@@ -114,4 +115,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
 
     serviceManager.addSingleton<IExportBase>(IExportBase, ExportBase);
     serviceManager.addSingleton<IExportUtil>(IExportUtil, ExportUtil);
+    serviceManager.addSingleton<NotebookPythonEnvironmentService>(
+        INotebookPythonEnvironmentService,
+        NotebookPythonEnvironmentService
+    );
 }

src/notebooks/serviceRegistry.web.ts

@@ -32,9 +32,10 @@ import { NotebookCellLanguageService } from './languages/cellLanguageService';
 import { EmptyNotebookCellLanguageService } from './languages/emptyNotebookCellLanguageService';
 import { NotebookCommandListener } from './notebookCommandListener';
 import { NotebookEditorProvider } from './notebookEditorProvider';
+import { NotebookPythonEnvironmentService } from './notebookEnvironmentService.web';
 import { CellOutputMimeTypeTracker } from './outputs/jupyterCellOutputMimeTypeTracker';
 import { NotebookTracebackFormatter } from './outputs/tracebackFormatter';
-import { INotebookEditorProvider } from './types';
+import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './types';
 
 export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) {
     registerControllerTypes(serviceManager, isDevMode);
@@ -87,4 +88,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
     serviceManager.addSingleton<IExportBase>(IExportBase, ExportBase);
     serviceManager.addSingleton<IFileConverter>(IFileConverter, FileConverter);
     serviceManager.addSingleton<IExportUtil>(IExportUtil, ExportUtil);
+    serviceManager.addSingleton<NotebookPythonEnvironmentService>(
+        INotebookPythonEnvironmentService,
+        NotebookPythonEnvironmentService
+    );
 }

src/notebooks/types.ts

@@ -1,8 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { NotebookDocument, NotebookEditor, Uri } from 'vscode';
+import { NotebookDocument, NotebookEditor, Uri, type Event } from 'vscode';
 import { Resource } from '../platform/common/types';
+import type { Environment } from '@vscode/python-extension';
 
 export interface IEmbedNotebookEditorProvider {
     findNotebookEditor(resource: Resource): NotebookEditor | undefined;
@@ -16,3 +17,9 @@ export interface INotebookEditorProvider {
     findAssociatedNotebookDocument(uri: Uri): NotebookDocument | undefined;
     registerEmbedNotebookProvider(provider: IEmbedNotebookEditorProvider): void;
 }
+
+export const INotebookPythonEnvironmentService = Symbol('INotebookPythonEnvironmentService');
+export interface INotebookPythonEnvironmentService {
+    onDidChangeEnvironment: Event<Uri>;
+    getPythonEnvironment(uri: Uri): Environment | undefined;
+}

src/platform/interpreter/helpers.ts

@@ -144,13 +144,19 @@ export function isCondaEnvironmentWithoutPython(interpreter?: { id: string }) {
     return env && getEnvironmentType(env) === EnvironmentType.Conda && !env.executable.uri;
 }
 
-export function getCachedEnvironment(interpreter?: { id: string }) {
+export function getCachedEnvironment(interpreter?: { id: string } | string) {
     if (!interpreter) {
         return;
     }
     if (!pythonApi) {
         throw new Error('Python API not initialized');
     }
+    if (typeof interpreter === 'string') {
+        return pythonApi.environments.known.find(
+            // eslint-disable-next-line local-rules/dont-use-fspath
+            (i) => i.id === interpreter || i.path === interpreter || i.executable.uri?.fsPath === interpreter
+        );
+    }
     return pythonApi.environments.known.find((i) => i.id === interpreter.id);
 }
 

src/platform/pythonEnvironments/info/index.ts

@@ -18,12 +18,10 @@ export enum EnvironmentType {
     VirtualEnvWrapper = 'VirtualEnvWrapper',
 }
 
-export type InterpreterId = string;
-
 /**
  * Details about a Python environment.
  */
 export interface PythonEnvironment {
-    id: InterpreterId;
+    id: string;
     uri: Uri;
 };