Jupyter API to get Env associated with Notebooks (#24771)

DonJayamanne · web-flow · commit b4e1ddb37bae · 2025-02-04T15:04:37.000-08:00
See microsoft/vscode-jupyter#15987 Should also fix microsoft/vscode-jupyter#16112 Should also avoid Pylance having to monitor notebook changes and then trying to figure out the Environment for a Notebook. previous discussion here #24358
diff --git a/pythonExtensionApi/src/main.ts b/pythonExtensionApi/src/main.ts
@@ -227,9 +227,9 @@ export type EnvironmentsChangeEvent = {
 
 export type ActiveEnvironmentPathChangeEvent = EnvironmentPath & {
     /**
-     * Workspace folder the environment changed for.
+     * Resource the environment changed for.
      */
-    readonly resource: WorkspaceFolder | undefined;
+    readonly resource: Resource | undefined;
 };
 
 /**
diff --git a/src/client/api.ts b/src/client/api.ts
@@ -15,7 +15,11 @@ import { IConfigurationService, Resource } from './common/types';
 import { getDebugpyLauncherArgs } from './debugger/extension/adapter/remoteLaunchers';
 import { IInterpreterService } from './interpreter/contracts';
 import { IServiceContainer, IServiceManager } from './ioc/types';
-import { JupyterExtensionIntegration } from './jupyter/jupyterIntegration';
+import {
+    JupyterExtensionIntegration,
+    JupyterExtensionPythonEnvironments,
+    JupyterPythonEnvironmentApi,
+} from './jupyter/jupyterIntegration';
 import { traceError } from './logging';
 import { IDiscoveryAPI } from './pythonEnvironments/base/locator';
 import { buildEnvironmentApi } from './environmentApi';
@@ -33,11 +37,16 @@ export function buildApi(
     const configurationService = serviceContainer.get<IConfigurationService>(IConfigurationService);
     const interpreterService = serviceContainer.get<IInterpreterService>(IInterpreterService);
     serviceManager.addSingleton<JupyterExtensionIntegration>(JupyterExtensionIntegration, JupyterExtensionIntegration);
+    serviceManager.addSingleton<JupyterExtensionPythonEnvironments>(
+        JupyterExtensionPythonEnvironments,
+        JupyterExtensionPythonEnvironments,
+    );
     serviceManager.addSingleton<TensorboardExtensionIntegration>(
         TensorboardExtensionIntegration,
         TensorboardExtensionIntegration,
     );
     const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
+    const jupyterPythonEnvApi = serviceContainer.get<JupyterPythonEnvironmentApi>(JupyterExtensionPythonEnvironments);
     const tensorboardIntegration = serviceContainer.get<TensorboardExtensionIntegration>(
         TensorboardExtensionIntegration,
     );
@@ -146,7 +155,7 @@ export function buildApi(
             stop: (client: BaseLanguageClient): Promise<void> => client.stop(),
             getTelemetryReporter: () => getTelemetryReporter(),
         },
-        environments: buildEnvironmentApi(discoveryApi, serviceContainer),
+        environments: buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi),
     };
 
     // In test environment return the DI Container.
diff --git a/src/client/api/types.ts b/src/client/api/types.ts
@@ -227,9 +227,9 @@ export type EnvironmentsChangeEvent = {
 
 export type ActiveEnvironmentPathChangeEvent = EnvironmentPath & {
     /**
-     * Workspace folder the environment changed for.
+     * Resource the environment changed for.
      */
-    readonly resource: WorkspaceFolder | undefined;
+    readonly resource: Resource | undefined;
 };
 
 /**
diff --git a/src/client/environmentApi.ts b/src/client/environmentApi.ts
@@ -33,6 +33,8 @@ import {
 } from './api/types';
 import { buildEnvironmentCreationApi } from './pythonEnvironments/creation/createEnvApi';
 import { EnvironmentKnownCache } from './environmentKnownCache';
+import type { JupyterPythonEnvironmentApi } from './jupyter/jupyterIntegration';
+import { noop } from './common/utils/misc';
 
 type ActiveEnvironmentChangeEvent = {
     resource: WorkspaceFolder | undefined;
@@ -115,6 +117,7 @@ function filterUsingVSCodeContext(e: PythonEnvInfo) {
 export function buildEnvironmentApi(
     discoveryApi: IDiscoveryAPI,
     serviceContainer: IServiceContainer,
+    jupyterPythonEnvsApi: JupyterPythonEnvironmentApi,
 ): PythonExtension['environments'] {
     const interpreterPathService = serviceContainer.get<IInterpreterPathService>(IInterpreterPathService);
     const configService = serviceContainer.get<IConfigurationService>(IConfigurationService);
@@ -146,6 +149,28 @@ export function buildEnvironmentApi(
             })
             .ignoreErrors();
     }
+
+    function getActiveEnvironmentPath(resource?: Resource) {
+        resource = resource && 'uri' in resource ? resource.uri : resource;
+        const jupyterEnv =
+            resource && jupyterPythonEnvsApi.getPythonEnvironment
+                ? jupyterPythonEnvsApi.getPythonEnvironment(resource)
+                : undefined;
+        if (jupyterEnv) {
+            traceVerbose('Python Environment returned from Jupyter', resource?.fsPath, jupyterEnv.id);
+            return {
+                id: jupyterEnv.id,
+                path: jupyterEnv.path,
+            };
+        }
+        const path = configService.getSettings(resource).pythonPath;
+        const id = path === 'python' ? 'DEFAULT_PYTHON' : getEnvID(path);
+        return {
+            id,
+            path,
+        };
+    }
+
     disposables.push(
         discoveryApi.onProgress((e) => {
             if (e.stage === ProgressReportStage.discoveryFinished) {
@@ -206,6 +231,16 @@ export function buildEnvironmentApi(
         }),
         onEnvironmentsChanged,
         onEnvironmentVariablesChanged,
+        jupyterPythonEnvsApi.onDidChangePythonEnvironment
+            ? jupyterPythonEnvsApi.onDidChangePythonEnvironment((e) => {
+                  const jupyterEnv = getActiveEnvironmentPath(e);
+                  onDidActiveInterpreterChangedEvent.fire({
+                      id: jupyterEnv.id,
+                      path: jupyterEnv.path,
+                      resource: e,
+                  });
+              }, undefined)
+            : { dispose: noop },
     );
     if (!knownCache!) {
         knownCache = initKnownCache();
@@ -223,13 +258,7 @@ export function buildEnvironmentApi(
         },
         getActiveEnvironmentPath(resource?: Resource) {
             sendApiTelemetry('getActiveEnvironmentPath');
-            resource = resource && 'uri' in resource ? resource.uri : resource;
-            const path = configService.getSettings(resource).pythonPath;
-            const id = path === 'python' ? 'DEFAULT_PYTHON' : getEnvID(path);
-            return {
-                id,
-                path,
-            };
+            return getActiveEnvironmentPath(resource);
         },
         updateActiveEnvironmentPath(env: Environment | EnvironmentPath | string, resource?: Resource): Promise<void> {
             sendApiTelemetry('updateActiveEnvironmentPath');
diff --git a/src/client/jupyter/jupyterIntegration.ts b/src/client/jupyter/jupyterIntegration.ts
@@ -1,12 +1,12 @@
 /* eslint-disable comma-dangle */
 
-/* eslint-disable implicit-arrow-linebreak */
+/* eslint-disable implicit-arrow-linebreak, max-classes-per-file */
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
 import { inject, injectable, named } from 'inversify';
 import { dirname } from 'path';
-import { Extension, Memento, Uri } from 'vscode';
+import { EventEmitter, Extension, Memento, Uri, workspace, Event } from 'vscode';
 import type { SemVer } from 'semver';
 import { IContextKeyManager, IWorkspaceService } from '../common/application/types';
 import { JUPYTER_EXTENSION_ID, PYLANCE_EXTENSION_ID } from '../common/constants';
@@ -23,6 +23,7 @@ import { PylanceApi } from '../activation/node/pylanceApi';
 import { ExtensionContextKey } from '../common/application/contextKeys';
 import { getDebugpyPath } from '../debugger/pythonDebugger';
 import type { Environment } from '../api/types';
+import { DisposableBase } from '../common/utils/resourceLifecycle';
 
 type PythonApiForJupyterExtension = {
     /**
@@ -170,3 +171,108 @@ export class JupyterExtensionIntegration {
         }
     }
 }
+
+export interface JupyterPythonEnvironmentApi {
+    /**
+     * 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;
+          };
+}
+
+@injectable()
+export class JupyterExtensionPythonEnvironments extends DisposableBase implements JupyterPythonEnvironmentApi {
+    private jupyterExtension?: JupyterPythonEnvironmentApi;
+
+    private readonly _onDidChangePythonEnvironment = this._register(new EventEmitter<Uri>());
+
+    public readonly onDidChangePythonEnvironment = this._onDidChangePythonEnvironment.event;
+
+    constructor(@inject(IExtensions) private readonly extensions: IExtensions) {
+        super();
+    }
+
+    public 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;
+          } {
+        if (!isJupyterResource(uri)) {
+            return undefined;
+        }
+        const api = this.getJupyterApi();
+        if (api?.getPythonEnvironment) {
+            return api.getPythonEnvironment(uri);
+        }
+        return undefined;
+    }
+
+    private getJupyterApi() {
+        if (!this.jupyterExtension) {
+            const ext = this.extensions.getExtension<JupyterPythonEnvironmentApi>(JUPYTER_EXTENSION_ID);
+            if (!ext) {
+                return undefined;
+            }
+            if (!ext.isActive) {
+                ext.activate().then(() => {
+                    this.hookupOnDidChangePythonEnvironment(ext.exports);
+                });
+                return undefined;
+            }
+            this.hookupOnDidChangePythonEnvironment(ext.exports);
+        }
+        return this.jupyterExtension;
+    }
+
+    private hookupOnDidChangePythonEnvironment(api: JupyterPythonEnvironmentApi) {
+        this.jupyterExtension = api;
+        if (api.onDidChangePythonEnvironment) {
+            this._register(
+                api.onDidChangePythonEnvironment(
+                    this._onDidChangePythonEnvironment.fire,
+                    this._onDidChangePythonEnvironment,
+                ),
+            );
+        }
+    }
+}
+
+function isJupyterResource(resource: Uri): boolean {
+    // Jupyter extension only deals with Notebooks and Interactive Windows.
+    return (
+        resource.fsPath.endsWith('.ipynb') ||
+        workspace.notebookDocuments.some((item) => item.uri.toString() === resource.toString())
+    );
+}
diff --git a/src/test/api.functional.test.ts b/src/test/api.functional.test.ts
@@ -19,6 +19,8 @@ import { ServiceManager } from '../client/ioc/serviceManager';
 import { IServiceContainer, IServiceManager } from '../client/ioc/types';
 import { IDiscoveryAPI } from '../client/pythonEnvironments/base/locator';
 import * as pythonDebugger from '../client/debugger/pythonDebugger';
+import { JupyterExtensionPythonEnvironments, JupyterPythonEnvironmentApi } from '../client/jupyter/jupyterIntegration';
+import { EventEmitter, Uri } from 'vscode';
 
 suite('Extension API', () => {
     const debuggerPath = path.join(EXTENSION_ROOT_DIR, 'python_files', 'lib', 'python', 'debugpy');
@@ -49,6 +51,14 @@ suite('Extension API', () => {
             instance(environmentVariablesProvider),
         );
         when(serviceContainer.get<IInterpreterService>(IInterpreterService)).thenReturn(instance(interpreterService));
+        const onDidChangePythonEnvironment = new EventEmitter<Uri>();
+        const jupyterApi: JupyterPythonEnvironmentApi = {
+            onDidChangePythonEnvironment: onDidChangePythonEnvironment.event,
+            getPythonEnvironment: (_uri: Uri) => undefined,
+        };
+        when(serviceContainer.get<JupyterPythonEnvironmentApi>(JupyterExtensionPythonEnvironments)).thenReturn(
+            jupyterApi,
+        );
         when(serviceContainer.get<IDisposableRegistry>(IDisposableRegistry)).thenReturn([]);
         getDebugpyPathStub = sinon.stub(pythonDebugger, 'getDebugpyPath');
         getDebugpyPathStub.resolves(debuggerPath);
diff --git a/src/test/environmentApi.unit.test.ts b/src/test/environmentApi.unit.test.ts
@@ -38,6 +38,7 @@ import {
     EnvironmentsChangeEvent,
     PythonExtension,
 } from '../client/api/types';
+import { JupyterPythonEnvironmentApi } from '../client/jupyter/jupyterIntegration';
 
 suite('Python Environment API', () => {
     const workspacePath = 'path/to/workspace';
@@ -80,7 +81,6 @@ suite('Python Environment API', () => {
         onDidChangeRefreshState = new EventEmitter();
         onDidChangeEnvironments = new EventEmitter();
         onDidChangeEnvironmentVariables = new EventEmitter();
-
         serviceContainer.setup((s) => s.get(IExtensions)).returns(() => extensions.object);
         serviceContainer.setup((s) => s.get(IInterpreterPathService)).returns(() => interpreterPathService.object);
         serviceContainer.setup((s) => s.get(IConfigurationService)).returns(() => configService.object);
@@ -94,8 +94,13 @@ suite('Python Environment API', () => {
         discoverAPI.setup((d) => d.onProgress).returns(() => onDidChangeRefreshState.event);
         discoverAPI.setup((d) => d.onChanged).returns(() => onDidChangeEnvironments.event);
         discoverAPI.setup((d) => d.getEnvs()).returns(() => []);
+        const onDidChangePythonEnvironment = new EventEmitter<Uri>();
+        const jupyterApi: JupyterPythonEnvironmentApi = {
+            onDidChangePythonEnvironment: onDidChangePythonEnvironment.event,
+            getPythonEnvironment: (_uri: Uri) => undefined,
+        };
 
-        environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object);
+        environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object, jupyterApi);
     });
 
     teardown(() => {
@@ -323,7 +328,12 @@ suite('Python Environment API', () => {
             },
         ];
         discoverAPI.setup((d) => d.getEnvs()).returns(() => envs);
-        environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object);
+        const onDidChangePythonEnvironment = new EventEmitter<Uri>();
+        const jupyterApi: JupyterPythonEnvironmentApi = {
+            onDidChangePythonEnvironment: onDidChangePythonEnvironment.event,
+            getPythonEnvironment: (_uri: Uri) => undefined,
+        };
+        environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object, jupyterApi);
         const actual = environmentApi.known;
         const actualEnvs = actual?.map((a) => (a as EnvironmentReference).internal);
         assert.deepEqual(
