feat: offer inline-script env setup as an unresolved-import quick fix

src/common/localize.ts

@@ -26,6 +26,12 @@ export namespace WorkbenchStrings {
 export namespace InlineScriptStrings {
     export const updateExtension = l10n.t('Update Extension');
 
+    export const setUpScriptEnvironment = l10n.t("Set up this script's Python environment");
+
+    export const saveFailedBeforeSetup = l10n.t(
+        'Could not save this script, so its environment was not set up. Save the file and try again.',
+    );
+
     export const updatePythonExtension = l10n.t(
         'The environment for this script was created. Update the Python extension for the full inline script experience.',
     );

src/features/inlineScript/setupCodeAction.ts

@@ -0,0 +1,121 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import {
+    CancellationToken,
+    CodeAction,
+    CodeActionContext,
+    CodeActionKind,
+    CodeActionProvider,
+    Diagnostic,
+    Disposable,
+    languages,
+    Range,
+    TextDocument,
+} from 'vscode';
+import { readInlineScriptMetadata, sliceHeaderBytes } from '../../common/inlineScript/metadata';
+import { getInlineScriptRoutingKey, InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry';
+import { InlineScriptStrings } from '../../common/localize';
+import { isInlineScriptsFeatureEnabled } from '../../helpers';
+
+/**
+ * Diagnostic codes meaning "this import did not resolve", lowercased for comparison.
+ *
+ * `reportMissingModuleSource` is included deliberately, unlike in Pylance's own
+ * `isMissingImportDiagnostic`: a stub without source means the package is not installed, which
+ * setting the script's environment up fixes.
+ */
+const UNRESOLVED_IMPORT_DIAGNOSTIC_CODES: ReadonlySet<string> = new Set([
+    // Pyright / Pylance / basedpyright.
+    'reportmissingimports',
+    'reportmissingmodulesource',
+    // Ty.
+    'unresolved-import',
+    'possibly-missing-import',
+    // Pyrefly.
+    'missing-import',
+    'missing-source',
+    'missing-source-for-stubs',
+    // mypy, via ms-python.mypy-type-checker.
+    'import-not-found',
+    'import-untyped',
+]);
+
+function normalizeDiagnosticCode(code: Diagnostic['code']): string | undefined {
+    if (code === undefined || code === null) {
+        return undefined;
+    }
+    const value = typeof code === 'object' ? code.value : code;
+    return typeof value === 'string' || typeof value === 'number' ? String(value).toLowerCase() : undefined;
+}
+
+/**
+ * Whether `diagnostic` reports an import that could not be resolved. Matches on `code`, never on
+ * `source`: Pyrefly-backed Pylance reports its source as the literal string `pylance + pyrefly`.
+ */
+export function isUnresolvedImportDiagnostic(diagnostic: Diagnostic): boolean {
+    const code = normalizeDiagnosticCode(diagnostic.code);
+    return code !== undefined && UNRESOLVED_IMPORT_DIAGNOSTIC_CODES.has(code);
+}
+
+/**
+ * Offers "Set up this script's Python environment" as a quick fix on an unresolved import in a `.py`
+ * file that declares a PEP 723 `# /// script` block and has no inline-script environment yet.
+ *
+ * Complements the CodeLens, which is hidden while the document is dirty — the moment a user has just
+ * typed the import that does not resolve. This provider parses the in-memory buffer instead.
+ *
+ * `diagnostics` and `isPreferred` are both left unset: setup installs the block's declared
+ * dependencies verbatim and may not resolve the import at all, so the action must not claim to fix
+ * the diagnostic or pre-empt a real import fix.
+ */
+export class InlineScriptSetupCodeActionProvider implements CodeActionProvider {
+    constructor(
+        private readonly routing: InlineScriptRoutingRegistry,
+        private readonly setupCommand: string,
+    ) {}
+
+    /** Gates run cheapest-first, and before any parsing: VS Code may call this on every cursor move. */
+    public provideCodeActions(
+        document: TextDocument,
+        _range: Range,
+        context: CodeActionContext,
+        _token: CancellationToken,
+    ): CodeAction[] {
+        if (!isInlineScriptsFeatureEnabled()) {
+            return [];
+        }
+        if (!context.diagnostics.some(isUnresolvedImportDiagnostic)) {
+            return [];
+        }
+        const uri = document.uri;
+        if (!getInlineScriptRoutingKey(uri)) {
+            return [];
+        }
+        if (this.routing.shouldRoute(uri)) {
+            return [];
+        }
+        if (!readInlineScriptMetadata(sliceHeaderBytes(document.getText()), uri.fsPath)) {
+            return [];
+        }
+        const action = new CodeAction(InlineScriptStrings.setUpScriptEnvironment, CodeActionKind.QuickFix);
+        action.command = {
+            title: InlineScriptStrings.setUpScriptEnvironment,
+            command: this.setupCommand,
+            arguments: [uri],
+        };
+        return [action];
+    }
+}
+
+/** Register the inline-script quick fix for local `.py` files. */
+export function registerInlineScriptSetupCodeAction(
+    routing: InlineScriptRoutingRegistry,
+    setupCommand: string,
+): Disposable {
+    return languages.registerCodeActionsProvider(
+        { scheme: 'file', language: 'python' },
+        new InlineScriptSetupCodeActionProvider(routing, setupCommand),
+        { providedCodeActionKinds: [CodeActionKind.QuickFix] },
+    );
+}

src/features/inlineScript/setupEnvironment.ts

@@ -1,12 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { commands, Disposable, l10n, QuickPickItem, Uri, window } from 'vscode';
+import { commands, Disposable, l10n, QuickPickItem, TextDocument, Uri, window } from 'vscode';
 import { PythonEnvironment } from '../../api';
 import { INLINE_SCRIPT_MANAGER_ID } from '../../common/constants';
 import { readInlineScriptMetadataFromFile } from '../../common/inlineScript/metadata';
 import { InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry';
-import { traceError, traceInfo } from '../../common/logging';
+import { InlineScriptStrings } from '../../common/localize';
+import { traceError, traceInfo, traceVerbose } from '../../common/logging';
 import { normalizePath } from '../../common/utils/pathUtils';
 import {
     showErrorMessage,
@@ -18,6 +19,7 @@ import { asRelativePath, findFiles, getOpenTextDocuments } from '../../common/wo
 import { EnvironmentManagers } from '../../internal.api';
 import { registerInlineScriptCodeLens } from './codeLens';
 import { promptUpdateExtensionsForInlineScripts } from './extensionVersionCheck';
+import { registerInlineScriptSetupCodeAction } from './setupCodeAction';
 
 /**
  * Hidden command invoked by the inline-script CodeLens to set up the environment for one script.
@@ -86,20 +88,50 @@ async function seedRoutingMetadataForClosedScript(scriptUri: Uri, routing: Inlin
     if (routing.getMetadata(scriptUri)) {
         return;
     }
+    if (findOpenDocument(scriptUri)) {
+        return;
+    }
+    const metadata = await readInlineScriptMetadataFromFile(scriptUri);
+    if (metadata && !routing.getMetadata(scriptUri)) {
+        routing.setMetadata(scriptUri, metadata);
+    }
+}
+
+/** The open text document backing `scriptUri`, if the user has it open. */
+function findOpenDocument(scriptUri: Uri): TextDocument | undefined {
     const scriptPath = normalizePath(scriptUri.fsPath);
-    const isOpen = getOpenTextDocuments().some(
+    return getOpenTextDocuments().find(
         (document) => document.uri.scheme === 'file' && normalizePath(document.uri.fsPath) === scriptPath,
     );
-    if (isOpen) {
-        return;
+}
+
+/**
+ * Save `scriptUri` if it is open with unsaved changes, so setup reads what the user actually sees.
+ *
+ * Returns `false` when the document could not be saved; setup must not run in that case.
+ */
+async function saveScriptBeforeSetup(scriptUri: Uri, routing: InlineScriptRoutingRegistry): Promise<boolean> {
+    const document = findOpenDocument(scriptUri);
+    if (!document?.isDirty) {
+        return true;
+    }
+    if (!(await document.save())) {
+        traceError(`Could not save ${scriptUri.fsPath} before setting up its inline-script environment.`);
+        return false;
     }
+    traceVerbose(`Saved ${scriptUri.fsPath} before setting up its inline-script environment.`);
+    // Seeding here is load-bearing: without it a just-typed block goes from no metadata to an
+    // identity while `create` runs, which `setUpInlineScriptEnvironment` reads as a concurrent edit
+    // and silently skips the association.
     const metadata = await readInlineScriptMetadataFromFile(scriptUri);
-    if (metadata && !routing.getMetadata(scriptUri)) {
+    if (metadata) {
         routing.setMetadata(scriptUri, metadata);
     }
+    return true;
 }
 
-function setupInlineScriptEnvironmentHandler(
+/** Handler for the single-file setup command, shared by the CodeLens and the quick fix. */
+export function setupInlineScriptEnvironmentHandler(
     em: EnvironmentManagers,
     routing: InlineScriptRoutingRegistry,
 ): (scriptUri?: Uri) => Promise<void> {
@@ -112,21 +144,31 @@ function setupInlineScriptEnvironmentHandler(
             showErrorMessage(l10n.t('The inline script environment manager is not available yet. Try again shortly.'));
             return;
         }
+        if (!(await saveScriptBeforeSetup(uri, routing))) {
+            showErrorMessage(InlineScriptStrings.saveFailedBeforeSetup);
+            return;
+        }
+        let environment: PythonEnvironment | undefined;
         try {
-            const environment = await setUpInlineScriptEnvironment(uri, em, routing);
-            if (!environment) {
-                notifyInlineScriptSetupOutcome(uri, routing);
-                return;
-            }
-            await promptUpdateExtensionsForInlineScripts();
+            environment = 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.',
                 ),
             );
+            return;
+        }
+        if (!environment) {
+            notifyInlineScriptSetupOutcome(uri, routing);
+            return;
         }
+        // Kept out of the try: the environment is already set up, so a failure in this follow-up
+        // must not be reported to the user as a setup failure.
+        await promptUpdateExtensionsForInlineScripts().catch((error) =>
+            traceError('Failed to check companion extension versions for inline scripts:', error),
+        );
     };
 }
 
@@ -314,13 +356,15 @@ async function filterInlineScriptFiles(files: readonly Uri[]): Promise<Uri[]> {
 }
 
 /**
- * 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 single-file setup command is invoked by
- * the CodeLens and stays out of `package.json`; the bulk command is palette-gated behind the flag.
+ * Register the inline-script user-facing surfaces (the CodeLens, the quick fix, and their setup
+ * commands). Only called when the PEP 723 inline-script feature flag is enabled. The single-file
+ * setup command is invoked by both surfaces and stays out of `package.json`; the bulk command is
+ * palette-gated behind the flag.
  */
 export function registerInlineScriptUx(em: EnvironmentManagers, routing: InlineScriptRoutingRegistry): Disposable[] {
     return [
         registerInlineScriptCodeLens(routing, SETUP_INLINE_SCRIPT_ENV_COMMAND),
+        registerInlineScriptSetupCodeAction(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),