diff --git a/src/api.proposed.notebookEnvironment.ts b/src/api.proposed.notebookEnvironment.ts new file mode 100644 index 00000000000..68fa08715d4 --- /dev/null +++ b/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; + /** + * 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; + }; + } +} diff --git a/src/extension.node.ts b/src/extension.node.ts index 4e519ab38f3..48fec4e8e5a 100644 --- a/src/extension.node.ts +++ b/src/extension.node.ts @@ -127,7 +127,10 @@ export async function activate(context: IExtensionContext): Promise Promise.resolve(undefined), onDidStart: () => ({ dispose: noop }) - } + }, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + onDidChangePythonEnvironment: undefined as any, + getPythonEnvironment: () => undefined }; } } diff --git a/src/extension.web.ts b/src/extension.web.ts index 9ffca564c40..d46fde12a15 100644 --- a/src/extension.web.ts +++ b/src/extension.web.ts @@ -117,10 +117,10 @@ export async function activate(context: IExtensionContext): Promise { 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 }; } } diff --git a/src/notebooks/notebookEnvironmentService.node.ts b/src/notebooks/notebookEnvironmentService.node.ts new file mode 100644 index 00000000000..a2c1157ae67 --- /dev/null +++ b/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()); + public readonly onDidChangeEnvironment = this._onDidChangeEnvironment.event; + + private readonly notebookWithRemoteKernelsToMonitor = new WeakSet(); + private readonly notebookPythonEnvironments = new WeakMap(); + 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 + ) { + 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 { + // 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 + )}` + ); + } + } +} diff --git a/src/notebooks/notebookEnvironmentService.web.ts b/src/notebooks/notebookEnvironmentService.web.ts new file mode 100644 index 00000000000..f718c3caefb --- /dev/null +++ b/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()); + public readonly onDidChangeEnvironment = this._onDidChangeEnvironment.event; + + public getPythonEnvironment(_: Uri): Environment | undefined { + return undefined; + } +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index cd1e980eaad..b3b6687609c 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/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, ExportBase); serviceManager.addSingleton(IExportUtil, ExportUtil); + serviceManager.addSingleton( + INotebookPythonEnvironmentService, + NotebookPythonEnvironmentService + ); } diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index dc9fe47cd1f..2833ce5ea5f 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/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, ExportBase); serviceManager.addSingleton(IFileConverter, FileConverter); serviceManager.addSingleton(IExportUtil, ExportUtil); + serviceManager.addSingleton( + INotebookPythonEnvironmentService, + NotebookPythonEnvironmentService + ); } diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 9ababf7d8ef..706a2345f13 100644 --- a/src/notebooks/types.ts +++ b/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; + getPythonEnvironment(uri: Uri): Environment | undefined; +} diff --git a/src/platform/interpreter/helpers.ts b/src/platform/interpreter/helpers.ts index 39ab0be7694..f2d5bff5995 100644 --- a/src/platform/interpreter/helpers.ts +++ b/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); } diff --git a/src/platform/pythonEnvironments/info/index.ts b/src/platform/pythonEnvironments/info/index.ts index 5f4b64783d1..221f597419c 100644 --- a/src/platform/pythonEnvironments/info/index.ts +++ b/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; }; diff --git a/src/standalone/api/index.ts b/src/standalone/api/index.ts index a7dec41f5bc..0820d49436f 100644 --- a/src/standalone/api/index.ts +++ b/src/standalone/api/index.ts @@ -18,6 +18,7 @@ import { openNotebook, registerRemoteServerProvider } from './unstable'; +import { INotebookPythonEnvironmentService } from '../../notebooks/types'; /* * Do not introduce any breaking changes to this API. @@ -34,6 +35,7 @@ export function buildApi( context: IExtensionContext ): IExtensionApi { const extensions = serviceContainer.get(IExtensions); + const envApi = serviceContainer.get(INotebookPythonEnvironmentService); const api: IExtensionApi = { // 'ready' will propagate the exception, but we must log it here first. ready: getReady(ready), @@ -56,6 +58,17 @@ export function buildApi( }, get kernels() { return getKernelsApi(extensions.determineExtensionFromCallStack().extensionId); + }, + // Only for use By Python, hence this is proposed API and can change anytime. + // Do not add this to Kernels or other namespaces, as this is only for Python. + onDidChangePythonEnvironment: envApi.onDidChangeEnvironment, + // Only for use By Python, hence this is proposed API and can change anytime. + // Do not add this to Kernels or other namespaces, as this is only for Python. + getPythonEnvironment(uri: Uri): EnvironmentPath | undefined { + // This is a proposed API that is only used by 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). + return envApi.getPythonEnvironment(uri); } };