Jupyter API to get Env associated with Notebooks

Author: DonJayamanne

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;
 };
 
 /**

src/client/api.ts

@@ -15,7 +15,7 @@ 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, type JupyterPythonEnvironmentApi } from './jupyter/jupyterIntegration';
 import { traceError } from './logging';
 import { IDiscoveryAPI } from './pythonEnvironments/base/locator';
 import { buildEnvironmentApi } from './environmentApi';
@@ -33,11 +33,13 @@ 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 +148,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.

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;
 };
 
 /**

src/client/environmentApi.ts

@@ -33,6 +33,7 @@ import {
 } from './api/types';
 import { buildEnvironmentCreationApi } from './pythonEnvironments/creation/createEnvApi';
 import { EnvironmentKnownCache } from './environmentKnownCache';
+import type { JupyterPythonEnvironmentApi } from './jupyter/jupyterIntegration';
 
 type ActiveEnvironmentChangeEvent = {
     resource: WorkspaceFolder | undefined;
@@ -115,6 +116,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);
@@ -206,6 +208,14 @@ export function buildEnvironmentApi(
         }),
         onEnvironmentsChanged,
         onEnvironmentVariablesChanged,
+        jupyterPythonEnvsApi.onDidChangePythonEnvironment((e) => {
+            const jupyterEnv = environmentApi.getActiveEnvironmentPath(e);
+            onDidActiveInterpreterChangedEvent.fire({
+                id: jupyterEnv.id,
+                path: jupyterEnv.path,
+                resource: e,
+            });
+        }, undefined),
     );
     if (!knownCache!) {
         knownCache = initKnownCache();
@@ -224,6 +234,14 @@ export function buildEnvironmentApi(
         getActiveEnvironmentPath(resource?: Resource) {
             sendApiTelemetry('getActiveEnvironmentPath');
             resource = resource && 'uri' in resource ? resource.uri : resource;
+            const jupyterEnv = resource ? 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 {

src/client/jupyter/jupyterIntegration.ts

@@ -6,7 +6,7 @@
 
 import { inject, injectable, named } from 'inversify';
 import { dirname } from 'path';
-import { Extension, Memento, Uri } from 'vscode';
+import { EventEmitter, Extension, Memento, Uri, workspace, type 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,83 @@ 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;
+            };
+}
+
+// eslint-disable-next-line max-classes-per-file
+
+@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;
+    } {
+        return isJupyterResource(uri) ? this.getJupyterApi()?.getPythonEnvironment(uri) : undefined;
+    }
+
+    private getJupyterApi(){
+        if (!this.jupyterExtension) {
+            const api = this.extensions.getExtension<JupyterPythonEnvironmentApi>(JUPYTER_EXTENSION_ID)?.exports;
+            if (!api) {
+                return undefined;
+            }
+            this.jupyterExtension = api;
+            this._register(api.onDidChangePythonEnvironment(this._onDidChangePythonEnvironment.fire, this._onDidChangePythonEnvironment));
+        }
+        return this.jupyterExtension;
+    }
+}
+
+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());
+}

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(