diff --git a/src/common/workspace.apis.ts b/src/common/workspace.apis.ts index d571bcf2..d8f48d6c 100644 --- a/src/common/workspace.apis.ts +++ b/src/common/workspace.apis.ts @@ -50,6 +50,10 @@ export function findFiles( return workspace.findFiles(include, exclude, maxResults, token); } +export function asRelativePath(pathOrUri: string | Uri, includeWorkspaceFolder?: boolean): string { + return workspace.asRelativePath(pathOrUri, includeWorkspaceFolder); +} + export function createFileSystemWatcher( globPattern: GlobPattern, ignoreCreateEvents?: boolean, diff --git a/src/extension.ts b/src/extension.ts index a5b393aa..f55169d1 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -69,6 +69,7 @@ import { PythonEnvironmentManagers } from './features/envManagers'; import { EnvVarManager, PythonEnvVariableManager } from './features/execution/envVariableManager'; import { latchInlineScriptFeatureActivation } from './features/inlineScript/activation'; import { InlineScriptLazyDetector } from './features/inlineScript/lazyDetector'; +import { registerInlineScriptUx } from './features/inlineScript/setupEnvironment'; import { applyInitialEnvironmentSelection, registerInterpreterSettingsChangeListener, @@ -411,6 +412,7 @@ export async function activate(context: ExtensionContext): Promise { return runInTerminalCommand(item, api, terminalManager); }), diff --git a/src/features/inlineScript/codeLens.ts b/src/features/inlineScript/codeLens.ts new file mode 100644 index 00000000..f605134c --- /dev/null +++ b/src/features/inlineScript/codeLens.ts @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { + CancellationToken, + CodeLens, + CodeLensProvider, + Disposable, + EventEmitter, + l10n, + languages, + Range, + TextDocument, +} from 'vscode'; +import { InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry'; + +/** + * Shows a single "Set up environment for this script" CodeLens above a `.py` file's PEP 723 + * `# /// script` block, but only when the file has saved inline metadata that is not currently + * backed by a validated inline-script environment. + * + * The provider is a pure observer of {@link InlineScriptRoutingRegistry}: + * - `getMetadata` returns the last saved metadata (the detector clears it while the metadata + * region is dirty), so the lens tracks *saved* metadata and disappears while it is being edited. + * - `shouldRoute` is true once a validated association matching the current metadata exists, so the + * lens hides after setup and reappears if a metadata change later invalidates that association. + */ +export class InlineScriptCodeLensProvider implements CodeLensProvider, Disposable { + private readonly _onDidChangeCodeLenses = new EventEmitter(); + public readonly onDidChangeCodeLenses = this._onDidChangeCodeLenses.event; + private readonly subscriptions: Disposable[] = []; + + constructor( + private readonly routing: InlineScriptRoutingRegistry, + private readonly setupCommand: string, + ) { + this.subscriptions.push( + this.routing.onDidChangeRouteability(() => this._onDidChangeCodeLenses.fire()), + // Only metadata arriving or changing can add/replace a lens; a scan that finds no metadata + // (the common case for ordinary .py files) needs no refresh. Hiding a lens for an + // edited/removed block is handled by VS Code re-querying on the document change itself. + this.routing.onDidChangeMetadata((e) => { + if (e.metadata !== undefined) { + this._onDidChangeCodeLenses.fire(); + } + }), + ); + } + + public provideCodeLenses(document: TextDocument, _token: CancellationToken): CodeLens[] { + if (document.isDirty) { + // The association is validated against the saved file (the manager refuses to validate a + // dirty document), so only offer setup for a clean document. This also avoids anchoring the + // lens at a stale offset if the block moved on an unsaved edit. + return []; + } + const uri = document.uri; + const metadata = this.routing.getMetadata(uri); + if (!metadata) { + // No saved PEP 723 metadata (or it is currently being edited). + return []; + } + if (this.routing.shouldRoute(uri)) { + // A validated inline-script environment matching the current metadata already exists. + return []; + } + const offset = metadata.sourceRange?.start ?? metadata.range.start; + const position = document.positionAt(offset); + const range = new Range(position, position); + return [ + new CodeLens(range, { + title: l10n.t('Set up environment for this script'), + command: this.setupCommand, + arguments: [uri], + }), + ]; + } + + public dispose(): void { + this.subscriptions.forEach((s) => s.dispose()); + this.subscriptions.length = 0; + this._onDidChangeCodeLenses.dispose(); + } +} + +/** + * Register the inline-script CodeLens provider for local `.py` files. Only called when the PEP 723 + * inline-script feature flag is enabled, so it is a no-op for everyone else. + */ +export function registerInlineScriptCodeLens(routing: InlineScriptRoutingRegistry, setupCommand: string): Disposable { + const provider = new InlineScriptCodeLensProvider(routing, setupCommand); + const registration = languages.registerCodeLensProvider({ scheme: 'file', language: 'python' }, provider); + return new Disposable(() => { + registration.dispose(); + provider.dispose(); + }); +} diff --git a/src/features/inlineScript/setupEnvironment.ts b/src/features/inlineScript/setupEnvironment.ts new file mode 100644 index 00000000..e071536e --- /dev/null +++ b/src/features/inlineScript/setupEnvironment.ts @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { commands, Disposable, l10n, QuickPickItem, Uri, window } from 'vscode'; +import { PythonEnvironment } from '../../api'; +import { INLINE_SCRIPT_MANAGER_ID } from '../../common/constants'; +import { InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry'; +import { readInlineScriptMetadataFromFile } from '../../common/inlineScript/metadata'; +import { traceError, traceInfo } from '../../common/logging'; +import { showErrorMessage, showInformationMessage, showQuickPickWithButtons } from '../../common/window.apis'; +import { asRelativePath, findFiles } from '../../common/workspace.apis'; +import { EnvironmentManagers } from '../../internal.api'; +import { registerInlineScriptCodeLens } from './codeLens'; + +/** + * Hidden command invoked by the inline-script CodeLens to set up the environment for one script. + * Intentionally not contributed in `package.json` while the feature is behind the internal flag. + */ +export const SETUP_INLINE_SCRIPT_ENV_COMMAND = 'python-envs.setupInlineScriptEnv'; + +/** + * Hidden command that scans the workspace and sets up environments for the selected inline-script + * files. Intentionally not contributed in `package.json` while the feature is behind the internal flag. + */ +export const SETUP_INLINE_SCRIPT_ENVS_COMMAND = 'python-envs.setupInlineScriptEnvs'; + +/** Upper bound on the number of `.py` files the bulk command scans. */ +const MAX_INLINE_SCRIPT_FILES = 1000; + +/** How many candidate files' PEP 723 headers are read concurrently during the bulk scan. */ +const METADATA_READ_CONCURRENCY = 20; + +/** + * Create or reuse the inline-script environment for `scriptUri` and make it the file's environment. + * + * Reuses the pipeline built by earlier PEP 723 PRs: + * - `manager.create` builds or reuses the cached environment from the script's PEP 723 metadata + * (selecting or, with consent, installing a compatible base interpreter); + * - `setEnvironment` persists the association, registers the exact script project, and publishes the + * per-file environment change so routing picks up the inline environment. + * + * Returns the environment on success, or `undefined` when creation produced none (the manager has + * already surfaced the reason — a declined install, no compatible Python, or a cancelled/failed + * build — and emitted telemetry). + */ +export async function setUpInlineScriptEnvironment( + scriptUri: Uri, + em: EnvironmentManagers, + routing: InlineScriptRoutingRegistry, +): Promise { + const manager = em.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID); + if (!manager) { + traceError('Inline-script setup requested but the inline-script environment manager is not registered.'); + return undefined; + } + const metadataIdentityBeforeCreate = routing.getMetadataIdentity(scriptUri); + const environment = await manager.create(scriptUri, undefined); + if (!environment) { + return undefined; + } + if (routing.getMetadataIdentity(scriptUri) !== metadataIdentityBeforeCreate) { + // The script's saved metadata changed while the environment was being built, so this + // environment was built for stale metadata. Skip associating it rather than overwrite a newer + // setup; the current metadata's CodeLens stays so the user can run setup again. + traceInfo(`Inline-script metadata for ${scriptUri.fsPath} changed during setup; skipping association.`); + return undefined; + } + await em.setEnvironment(scriptUri, environment); + return environment; +} + +function setupInlineScriptEnvironmentHandler( + em: EnvironmentManagers, + routing: InlineScriptRoutingRegistry, +): (scriptUri?: Uri) => Promise { + return async (scriptUri?: Uri): Promise => { + const uri = scriptUri ?? window.activeTextEditor?.document.uri; + if (!uri || uri.scheme !== 'file') { + return; + } + if (!em.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)) { + showErrorMessage(l10n.t('The inline script environment manager is not available yet. Try again shortly.')); + return; + } + try { + await setUpInlineScriptEnvironment(uri, em, routing); + } catch (error) { + traceError(`Failed to set up the inline-script environment for ${uri.fsPath}:`, error); + showErrorMessage( + l10n.t( + 'Failed to set up the environment for this script. See the Python Environments output for details.', + ), + ); + } + }; +} + +interface InlineScriptQuickPickItem extends QuickPickItem { + readonly uri: Uri; + readonly configured: boolean; +} + +export async function setUpInlineScriptEnvironmentsInWorkspace( + em: EnvironmentManagers, + routing: InlineScriptRoutingRegistry, +): Promise { + if (!em.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)) { + showErrorMessage(l10n.t('The inline script environment manager is not available yet. Try again shortly.')); + return; + } + const files = await findFiles('**/*.py', '{**/.venv/**,**/node_modules/**}', MAX_INLINE_SCRIPT_FILES); + if (!files || files.length === 0) { + showInformationMessage(l10n.t('No Python files were found in the workspace.')); + return; + } + const candidates = await filterInlineScriptFiles(files); + if (candidates.length === 0) { + showInformationMessage( + l10n.t('No Python files with PEP 723 inline script metadata were found in the workspace.'), + ); + return; + } + const items: InlineScriptQuickPickItem[] = candidates + .map((uri) => { + const configured = routing.shouldRoute(uri); + return { + label: asRelativePath(uri), + description: configured ? l10n.t('environment already set up') : undefined, + uri, + configured, + }; + }) + .sort((a, b) => a.label.localeCompare(b.label)); + // Pre-select only scripts that do not already have a validated inline environment, so that + // accepting the picker never silently rebuilds or replaces an already-configured script's + // environment (which could otherwise move it onto a newer base interpreter). + const preselected = items.filter((item) => !item.configured); + const selection = await showQuickPickWithButtons(items, { + canPickMany: true, + ignoreFocusOut: true, + title: l10n.t('Set Up Environments for Inline Script Files'), + placeHolder: l10n.t('Select the scripts to set up environments for'), + selected: preselected, + }); + const picks = Array.isArray(selection) ? selection : selection ? [selection] : []; + if (picks.length === 0) { + return; + } + let succeeded = 0; + for (const pick of picks) { + try { + if (await setUpInlineScriptEnvironment(pick.uri, em, routing)) { + succeeded += 1; + } + } catch (error) { + traceError(`Failed to set up the inline-script environment for ${pick.uri.fsPath}:`, error); + } + } + traceInfo(`Inline-script bulk setup: created or reused ${succeeded} of ${picks.length} environment(s).`); + showInformationMessage(l10n.t('Set up {0} of {1} selected inline script environment(s).', succeeded, picks.length)); +} + +/** + * Read the head of each candidate file (bounded concurrency) and keep only those that declare a + * PEP 723 `# /// script` block. + */ +async function filterInlineScriptFiles(files: readonly Uri[]): Promise { + const candidates: Uri[] = []; + for (let index = 0; index < files.length; index += METADATA_READ_CONCURRENCY) { + const chunk = files.slice(index, index + METADATA_READ_CONCURRENCY); + const results = await Promise.all( + chunk.map(async (uri) => ((await readInlineScriptMetadataFromFile(uri)) ? uri : undefined)), + ); + for (const uri of results) { + if (uri) { + candidates.push(uri); + } + } + } + return candidates; +} + +/** + * Register the inline-script user-facing surfaces (the CodeLens and its setup commands). Only called + * when the PEP 723 inline-script feature flag is enabled; the commands are intentionally hidden from + * `package.json` for now. + */ +export function registerInlineScriptUx(em: EnvironmentManagers, routing: InlineScriptRoutingRegistry): Disposable[] { + return [ + registerInlineScriptCodeLens(routing, SETUP_INLINE_SCRIPT_ENV_COMMAND), + commands.registerCommand(SETUP_INLINE_SCRIPT_ENV_COMMAND, setupInlineScriptEnvironmentHandler(em, routing)), + commands.registerCommand(SETUP_INLINE_SCRIPT_ENVS_COMMAND, () => + setUpInlineScriptEnvironmentsInWorkspace(em, routing), + ), + ]; +} diff --git a/src/test/features/inlineScript/codeLens.unit.test.ts b/src/test/features/inlineScript/codeLens.unit.test.ts new file mode 100644 index 00000000..935af095 --- /dev/null +++ b/src/test/features/inlineScript/codeLens.unit.test.ts @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import assert from 'assert'; +import { Position, TextDocument, Uri } from 'vscode'; +import { InlineScriptMetadata } from '../../../common/inlineScript/metadata'; +import { InlineScriptRoutingRegistry } from '../../../common/inlineScript/routingRegistry'; +import { InlineScriptCodeLensProvider } from '../../../features/inlineScript/codeLens'; + +const SETUP_COMMAND = 'python-envs.setupInlineScriptEnv'; + +function makeMetadata(): InlineScriptMetadata { + return { + dependencies: ['requests'], + range: { start: 0, end: 24 }, + sourceRange: { start: 0, end: 24 }, + }; +} + +function makeDocument(uri: Uri, isDirty = false): TextDocument { + return { + uri, + isDirty, + positionAt: (offset: number) => new Position(0, offset), + } as unknown as TextDocument; +} + +suite('Inline script CodeLens provider', () => { + const scriptUri = Uri.file('/workspace/app.py'); + let routing: InlineScriptRoutingRegistry; + let provider: InlineScriptCodeLensProvider; + + setup(() => { + routing = new InlineScriptRoutingRegistry(); + provider = new InlineScriptCodeLensProvider(routing, SETUP_COMMAND); + }); + + teardown(() => { + provider.dispose(); + routing.dispose(); + }); + + test('shows no CodeLens when the file has no saved inline metadata', () => { + const lenses = provider.provideCodeLenses(makeDocument(scriptUri), {} as never); + assert.strictEqual(lenses.length, 0); + }); + + test('shows a setup CodeLens when metadata exists but no environment is associated', () => { + routing.setMetadata(scriptUri, makeMetadata()); + + const lenses = provider.provideCodeLenses(makeDocument(scriptUri), {} as never); + + assert.strictEqual(lenses.length, 1); + assert.strictEqual(lenses[0].command?.command, SETUP_COMMAND); + assert.deepStrictEqual(lenses[0].command?.arguments, [scriptUri]); + }); + + test('shows no CodeLens while the document has unsaved changes', () => { + routing.setMetadata(scriptUri, makeMetadata()); + + const lenses = provider.provideCodeLenses(makeDocument(scriptUri, true), {} as never); + + assert.strictEqual(lenses.length, 0); + }); + + test('hides the CodeLens once a validated association makes the script routeable', () => { + routing.setMetadata(scriptUri, makeMetadata()); + routing.setValidatedAssociation(scriptUri, true); + assert.strictEqual(routing.shouldRoute(scriptUri), true); + + const lenses = provider.provideCodeLenses(makeDocument(scriptUri), {} as never); + + assert.strictEqual(lenses.length, 0); + }); + + test('refreshes CodeLenses when routing state changes', () => { + let fireCount = 0; + const sub = provider.onDidChangeCodeLenses(() => (fireCount += 1)); + + routing.setMetadata(scriptUri, makeMetadata()); + routing.setValidatedAssociation(scriptUri, true); + + sub.dispose(); + assert.ok(fireCount >= 1, 'onDidChangeCodeLenses should fire when routing state changes'); + }); +}); diff --git a/src/test/features/inlineScript/setupEnvironment.unit.test.ts b/src/test/features/inlineScript/setupEnvironment.unit.test.ts new file mode 100644 index 00000000..176fe555 --- /dev/null +++ b/src/test/features/inlineScript/setupEnvironment.unit.test.ts @@ -0,0 +1,164 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import assert from 'assert'; +import * as sinon from 'sinon'; +import * as typemoq from 'typemoq'; +import { Uri } from 'vscode'; +import { PythonEnvironment } from '../../../api'; +import { INLINE_SCRIPT_MANAGER_ID } from '../../../common/constants'; +import { InlineScriptMetadata } from '../../../common/inlineScript/metadata'; +import * as metadataApi from '../../../common/inlineScript/metadata'; +import { InlineScriptRoutingRegistry } from '../../../common/inlineScript/routingRegistry'; +import * as winapi from '../../../common/window.apis'; +import * as wapi from '../../../common/workspace.apis'; +import { + setUpInlineScriptEnvironment, + setUpInlineScriptEnvironmentsInWorkspace, +} from '../../../features/inlineScript/setupEnvironment'; +import { EnvironmentManagers, InternalEnvironmentManager } from '../../../internal.api'; + +function makeEnv(): PythonEnvironment { + return { + envId: { id: 'env1', managerId: INLINE_SCRIPT_MANAGER_ID }, + name: 'inline', + version: '3.12.0', + environmentPath: Uri.file('/cache/env/python'), + displayName: 'inline', + displayPath: '/cache/env/python', + execInfo: { run: { executable: '/cache/env/python' } }, + sysPrefix: '/cache/env', + } as PythonEnvironment; +} + +function makeMetadata(dependencies: string[]): InlineScriptMetadata { + return { + dependencies, + range: { start: 0, end: 10 }, + sourceRange: { start: 0, end: 10 }, + }; +} + +suite('setUpInlineScriptEnvironment', () => { + const scriptUri = Uri.file('/workspace/app.py'); + let em: typemoq.IMock; + let manager: typemoq.IMock; + let routing: InlineScriptRoutingRegistry; + + setup(() => { + em = typemoq.Mock.ofType(); + manager = typemoq.Mock.ofType(); + routing = new InlineScriptRoutingRegistry(); + em.setup((m) => m.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)).returns(() => manager.object); + }); + + teardown(() => { + routing.dispose(); + sinon.restore(); + }); + + test('returns undefined and sets no environment when the inline manager is not registered', async () => { + em.reset(); + em.setup((m) => m.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)).returns(() => undefined); + + const result = await setUpInlineScriptEnvironment(scriptUri, em.object, routing); + + assert.strictEqual(result, undefined); + em.verify((m) => m.setEnvironment(typemoq.It.isAny(), typemoq.It.isAny()), typemoq.Times.never()); + }); + + test('does not set an environment when creation produces none', async () => { + manager.setup((m) => m.create(scriptUri, undefined)).returns(() => Promise.resolve(undefined)); + + const result = await setUpInlineScriptEnvironment(scriptUri, em.object, routing); + + assert.strictEqual(result, undefined); + em.verify((m) => m.setEnvironment(typemoq.It.isAny(), typemoq.It.isAny()), typemoq.Times.never()); + }); + + test('creates then sets the environment for the script', async () => { + const env = makeEnv(); + manager.setup((m) => m.create(scriptUri, undefined)).returns(() => Promise.resolve(env)); + em.setup((m) => m.setEnvironment(scriptUri, env)).returns(() => Promise.resolve()); + + const result = await setUpInlineScriptEnvironment(scriptUri, em.object, routing); + + assert.strictEqual(result, env); + em.verify((m) => m.setEnvironment(scriptUri, env), typemoq.Times.once()); + }); + + test('skips association when the script metadata changes during creation', async () => { + routing.setMetadata(scriptUri, makeMetadata(['a'])); + const env = makeEnv(); + manager + .setup((m) => m.create(scriptUri, undefined)) + .returns(async () => { + // Simulate the user editing + saving new dependencies while the environment is building. + routing.setMetadata(scriptUri, makeMetadata(['b'])); + return env; + }); + + const result = await setUpInlineScriptEnvironment(scriptUri, em.object, routing); + + assert.strictEqual(result, undefined); + em.verify((m) => m.setEnvironment(typemoq.It.isAny(), typemoq.It.isAny()), typemoq.Times.never()); + }); +}); + +suite('setUpInlineScriptEnvironmentsInWorkspace', () => { + const withMeta = Uri.file('/workspace/with_meta.py'); + const withoutMeta = Uri.file('/workspace/plain.py'); + let em: typemoq.IMock; + let manager: typemoq.IMock; + let routing: InlineScriptRoutingRegistry; + let findFilesStub: sinon.SinonStub; + let readMetadataStub: sinon.SinonStub; + let quickPickStub: sinon.SinonStub; + let infoStub: sinon.SinonStub; + + setup(() => { + em = typemoq.Mock.ofType(); + manager = typemoq.Mock.ofType(); + routing = new InlineScriptRoutingRegistry(); + em.setup((m) => m.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)).returns(() => manager.object); + + findFilesStub = sinon.stub(wapi, 'findFiles'); + sinon.stub(wapi, 'asRelativePath').callsFake((p) => (p instanceof Uri ? p.fsPath : String(p))); + readMetadataStub = sinon.stub(metadataApi, 'readInlineScriptMetadataFromFile'); + quickPickStub = sinon.stub(winapi, 'showQuickPickWithButtons'); + infoStub = sinon.stub(winapi, 'showInformationMessage'); + }); + + teardown(() => { + routing.dispose(); + sinon.restore(); + }); + + test('reports and sets up nothing when no files declare inline metadata', async () => { + findFilesStub.resolves([withoutMeta]); + readMetadataStub.resolves(undefined); + + await setUpInlineScriptEnvironmentsInWorkspace(em.object, routing); + + assert.ok(infoStub.calledOnce); + manager.verify((m) => m.create(typemoq.It.isAny(), typemoq.It.isAny()), typemoq.Times.never()); + }); + + test('only sets up the selected files that declare inline metadata', async () => { + findFilesStub.resolves([withMeta, withoutMeta]); + readMetadataStub.callsFake(async (uri: Uri) => + uri.fsPath === withMeta.fsPath ? makeMetadata(['requests']) : undefined, + ); + // Simulate the user accepting the pre-selected candidates. + quickPickStub.callsFake((items) => items); + const env = makeEnv(); + manager.setup((m) => m.create(withMeta, undefined)).returns(() => Promise.resolve(env)); + em.setup((m) => m.setEnvironment(withMeta, env)).returns(() => Promise.resolve()); + + await setUpInlineScriptEnvironmentsInWorkspace(em.object, routing); + + manager.verify((m) => m.create(withMeta, undefined), typemoq.Times.once()); + manager.verify((m) => m.create(withoutMeta, undefined), typemoq.Times.never()); + em.verify((m) => m.setEnvironment(withMeta, env), typemoq.Times.once()); + }); +});