From 862bd664d0609c32b3e9dbe30d81da1eab8b9203 Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Sun, 29 Dec 2024 23:42:20 +1100 Subject: [PATCH] Use Jupyter API to get Notebook/IW Python Env --- pythonExtensionApi/src/main.ts | 4 +- src/client/api.ts | 13 ++- src/client/api/types.ts | 4 +- src/client/environmentApi.ts | 43 ++++++++-- src/client/jupyter/jupyterIntegration.ts | 100 ++++++++++++++++++++++- src/test/api.functional.test.ts | 10 +++ src/test/environmentApi.unit.test.ts | 16 +++- 7 files changed, 172 insertions(+), 18 deletions(-) diff --git a/pythonExtensionApi/src/main.ts b/pythonExtensionApi/src/main.ts index 154ffbbd857a..2173245cbb28 100644 --- 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 index 899326647808..15fb4d688a89 100644 --- 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); const interpreterService = serviceContainer.get(IInterpreterService); serviceManager.addSingleton(JupyterExtensionIntegration, JupyterExtensionIntegration); + serviceManager.addSingleton( + JupyterExtensionPythonEnvironments, + JupyterExtensionPythonEnvironments, + ); serviceManager.addSingleton( TensorboardExtensionIntegration, TensorboardExtensionIntegration, ); const jupyterIntegration = serviceContainer.get(JupyterExtensionIntegration); + const jupyterPythonEnvApi = serviceContainer.get(JupyterExtensionPythonEnvironments); const tensorboardIntegration = serviceContainer.get( TensorboardExtensionIntegration, ); @@ -146,7 +155,7 @@ export function buildApi( stop: (client: BaseLanguageClient): Promise => 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 index 4e67334121fb..95556aacbd90 100644 --- 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 index 6c4b5cf94d92..558938d7d0b7 100644 --- 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); const configService = serviceContainer.get(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 { sendApiTelemetry('updateActiveEnvironmentPath'); diff --git a/src/client/jupyter/jupyterIntegration.ts b/src/client/jupyter/jupyterIntegration.ts index 69583b744da9..9c5952def555 100644 --- 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,98 @@ 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; + /** + * 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()); + + 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 api = this.extensions.getExtension(JUPYTER_EXTENSION_ID)?.exports; + if (!api) { + return undefined; + } + this.jupyterExtension = api; + if (api.onDidChangePythonEnvironment) { + 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()) + ); +} diff --git a/src/test/api.functional.test.ts b/src/test/api.functional.test.ts index eea0fb920b15..1149dcb7da9d 100644 --- 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)).thenReturn(instance(interpreterService)); + const onDidChangePythonEnvironment = new EventEmitter(); + const jupyterApi: JupyterPythonEnvironmentApi = { + onDidChangePythonEnvironment: onDidChangePythonEnvironment.event, + getPythonEnvironment: (_uri: Uri) => undefined, + }; + when(serviceContainer.get(JupyterExtensionPythonEnvironments)).thenReturn( + jupyterApi, + ); when(serviceContainer.get(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 index 012e1a0bfc69..2e5d13161f7b 100644 --- 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(); + 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(); + 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(