Add python-arch support (#1079)

Windows ARM64 jobs may need x64 Python independently of their Python
version pin. Add `python-arch` to export `UV_PYTHON_ARCH` for subsequent
steps and include the selected architecture in the uv and managed-Python
cache keys. The input takes precedence over an inherited environment
value, and unsupported uv versions fail before environment creation.

Requires a uv release with `UV_PYTHON_ARCH` support.

src/cache/restore-cache.ts

@@ -98,11 +98,14 @@ async function computeKeys(
     ? `-${encodeURIComponent(inputs.cacheSuffix)}`
     : "";
   const version = encodeURIComponent(pythonVersion ?? "unknown");
+  const pythonArchKey = inputs.pythonArch
+    ? `-python-${encodeURIComponent(inputs.pythonArch)}`
+    : "";
   const platform = await getPlatform();
   const osNameVersion = getOSNameVersion();
   const pruned = inputs.pruneCache ? "-pruned" : "";
   const python = inputs.cachePython ? "-py" : "";
-  return `setup-uv-${CACHE_VERSION}-${getArch()}-${platform}-${osNameVersion}-${version}${pruned}${python}${cacheDependencyPathHash}${suffix}`;
+  return `setup-uv-${CACHE_VERSION}-${getArch()}-${platform}-${osNameVersion}-${version}${pythonArchKey}${pruned}${python}${cacheDependencyPathHash}${suffix}`;
 }
 
 function handleMatchResult(

src/setup-uv.ts

@@ -16,6 +16,7 @@ import {
   getPlatform,
   type Platform,
 } from "./utils/platforms";
+import { setupPythonArch } from "./utils/python-arch";
 import { getPythonRuntimeId } from "./utils/python-runtime";
 import { resolveUvVersion } from "./version/resolve";
 
@@ -86,6 +87,11 @@ async function run(): Promise<void> {
       throw new Error(`Unsupported architecture: ${process.arch}`);
     }
     const setupResult = await setupUv(inputs, platform, arch);
+    const uvPath = path.join(
+      setupResult.uvDir,
+      process.platform === "win32" ? "uv.exe" : "uv",
+    );
+    await setupPythonArch(uvPath, inputs.pythonArch, inputs.exportPythonArch);
 
     addToolBinToPath(inputs);
     addUvToPathAndOutput(setupResult.uvDir);

src/utils/inputs.ts

@@ -24,6 +24,8 @@ export interface SetupInputs {
   version: string;
   versionFile: string;
   pythonVersion: string;
+  pythonArch: string;
+  exportPythonArch: boolean;
   activateEnvironment: boolean;
   noProject: boolean;
   venvPath: string;
@@ -54,6 +56,9 @@ export function loadInputs(): SetupInputs {
   const version = core.getInput("version");
   const versionFile = getVersionFile(workingDirectory);
   const pythonVersion = getPythonVersion(versionFile);
+  const pythonArchInput = core.getInput("python-arch");
+  const pythonArch = pythonArchInput || process.env.UV_PYTHON_ARCH || "";
+  const exportPythonArch = pythonArchInput !== "";
   const activateEnvironment = core.getBooleanInput("activate-environment");
   const noProject = core.getBooleanInput("no-project");
   const venvPath = getVenvPath(workingDirectory, activateEnvironment);
@@ -94,12 +99,14 @@ export function loadInputs(): SetupInputs {
     checksum,
     downloadFromAstralMirror,
     enableCache,
+    exportPythonArch,
     githubToken,
     ignoreEmptyWorkdir,
     ignoreNothingToCache,
     manifestFile,
     noProject,
     pruneCache,
+    pythonArch,
     pythonDir,
     pythonVersion,
     quiet,

src/utils/python-arch.ts

@@ -0,0 +1,75 @@
+import { execFile } from "node:child_process";
+import { promisify } from "node:util";
+import * as core from "@actions/core";
+import * as log from "./logging";
+
+const execFileAsync = promisify(execFile);
+const PROBE_ARCH = "setup-uv-probe";
+
+export async function setupPythonArch(
+  uvPath: string,
+  pythonArch: string,
+  exportPythonArch: boolean,
+): Promise<void> {
+  if (pythonArch === "") {
+    return;
+  }
+
+  // Older uv releases ignore unknown environment variables. An invalid architecture
+  // confirms that uv recognizes UV_PYTHON_ARCH without finding or downloading Python.
+  const probe = await queryPythonArch(uvPath, PROBE_ARCH);
+  if (probe.exitCode === 0) {
+    throw new Error(
+      "The installed version of uv does not support UV_PYTHON_ARCH. Select a newer uv version or use an architecture-qualified python-version.",
+    );
+  }
+  if (
+    !probe.stderr.includes(
+      `environment variable \`UV_PYTHON_ARCH\` with invalid value \`${PROBE_ARCH}\``,
+    )
+  ) {
+    throw new Error(
+      `Failed to check uv's support for UV_PYTHON_ARCH: ${probe.stderr.trim() || `uv exited with code ${probe.exitCode}`}`,
+    );
+  }
+
+  const selected = await queryPythonArch(uvPath, pythonArch);
+  if (selected.exitCode !== 0) {
+    throw new Error(
+      `Failed to set Python architecture to ${pythonArch}: ${selected.stderr.trim() || `uv exited with code ${selected.exitCode}`}`,
+    );
+  }
+
+  if (exportPythonArch) {
+    core.exportVariable("UV_PYTHON_ARCH", pythonArch);
+    log.info(`Set UV_PYTHON_ARCH to ${pythonArch}`);
+  }
+}
+
+async function queryPythonArch(
+  uvPath: string,
+  pythonArch: string,
+): Promise<{ exitCode: number; stderr: string }> {
+  try {
+    const { stderr } = await execFileAsync(
+      uvPath,
+      ["--no-config", "cache", "dir"],
+      {
+        encoding: "utf8",
+        env: { ...process.env, NO_COLOR: "1", UV_PYTHON_ARCH: pythonArch },
+      },
+    );
+    return { exitCode: 0, stderr };
+  } catch (error) {
+    if (
+      error instanceof Error &&
+      "code" in error &&
+      typeof error.code === "number" &&
+      "stderr" in error &&
+      typeof error.stderr === "string"
+    ) {
+      return { exitCode: error.code, stderr: error.stderr };
+    }
+    throw error;
+  }
+}