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.

__tests__/cache/restore-cache.test.ts

@@ -4,6 +4,9 @@ import { createSetupInputs } from "../helpers/setup-inputs";
 const mockRestoreCache = jest.fn();
 const mockSaveState = jest.fn();
 const mockSetOutput = jest.fn();
+const mockGetArch = jest.fn(() => "x86_64");
+const mockGetOSNameVersion = jest.fn(() => "ubuntu-24.04");
+const mockGetPlatform = jest.fn(async () => "unknown-linux-gnu");
 
 jest.unstable_mockModule("@actions/cache", () => ({
   restoreCache: mockRestoreCache,
@@ -24,9 +27,9 @@ jest.unstable_mockModule("../../src/utils/logging", () => ({
 }));
 
 jest.unstable_mockModule("../../src/utils/platforms", () => ({
-  getArch: jest.fn(() => "x86_64"),
-  getOSNameVersion: jest.fn(() => "ubuntu-24.04"),
-  getPlatform: jest.fn(async () => "unknown-linux-gnu"),
+  getArch: mockGetArch,
+  getOSNameVersion: mockGetOSNameVersion,
+  getPlatform: mockGetPlatform,
 }));
 
 const { restoreCache } = await import("../../src/cache/restore-cache");
@@ -39,6 +42,9 @@ function cacheKeyOutput(): string {
 
 beforeEach(() => {
   jest.clearAllMocks();
+  mockGetArch.mockReturnValue("x86_64");
+  mockGetOSNameVersion.mockReturnValue("ubuntu-24.04");
+  mockGetPlatform.mockResolvedValue("unknown-linux-gnu");
 });
 
 describe("restoreCache", () => {
@@ -71,4 +77,39 @@ describe("restoreCache", () => {
       "setup-uv-2-x86_64-unknown-linux-gnu-ubuntu-24.04-3.11-dependencyhash-tests-3.11",
     );
   });
+
+  it("includes the resolved Python architecture in the cache key", async () => {
+    const inputs = createSetupInputs({ pythonArch: "aarch64" });
+
+    await restoreCache(inputs, "3.14");
+
+    expect(cacheKeyOutput()).toContain("-3.14-python-aarch64-");
+  });
+
+  it.each(["aarch64", "x86_64"])(
+    "separates %s Python caches on a Windows ARM64 runner",
+    async (pythonArch) => {
+      mockGetArch.mockReturnValue("aarch64");
+      mockGetOSNameVersion.mockReturnValue("windows-11");
+      mockGetPlatform.mockResolvedValue("pc-windows-msvc");
+      const inputs = createSetupInputs({
+        cachePython: true,
+        pythonArch,
+        restoreCache: true,
+      });
+
+      await restoreCache(inputs, "3.14");
+
+      const cacheKey = `setup-uv-2-aarch64-pc-windows-msvc-windows-11-3.14-python-${pythonArch}-py-dependencyhash`;
+      expect(cacheKeyOutput()).toBe(cacheKey);
+      expect(mockRestoreCache).toHaveBeenCalledWith(
+        [inputs.cacheLocalPath?.path],
+        cacheKey,
+      );
+      expect(mockRestoreCache).toHaveBeenCalledWith(
+        [inputs.pythonDir],
+        `${cacheKey}-python`,
+      );
+    },
+  );
 });

__tests__/helpers/setup-inputs.ts

@@ -16,11 +16,13 @@ export function createSetupInputs(
     checksum: "",
     downloadFromAstralMirror: false,
     enableCache: true,
+    exportPythonArch: false,
     githubToken: "",
     ignoreEmptyWorkdir: false,
     ignoreNothingToCache: false,
     noProject: false,
     pruneCache: false,
+    pythonArch: "",
     pythonDir: "/tmp/uv-python-dir",
     pythonVersion: "",
     quiet: false,

__tests__/utils/inputs.test.ts

@@ -19,6 +19,7 @@ const ORIGINAL_RUNNER_ENVIRONMENT = process.env.RUNNER_ENVIRONMENT;
 const ORIGINAL_RUNNER_TEMP = process.env.RUNNER_TEMP;
 const ORIGINAL_UV_CACHE_DIR = process.env.UV_CACHE_DIR;
 const ORIGINAL_UV_PYTHON = process.env.UV_PYTHON;
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
 const ORIGINAL_UV_PYTHON_INSTALL_DIR = process.env.UV_PYTHON_INSTALL_DIR;
 
 const mockDebug = jest.fn();
@@ -62,6 +63,7 @@ function resetEnvironment(): void {
   delete process.env.RUNNER_TEMP;
   delete process.env.UV_CACHE_DIR;
   delete process.env.UV_PYTHON;
+  delete process.env.UV_PYTHON_ARCH;
   delete process.env.UV_PYTHON_INSTALL_DIR;
 }
 
@@ -77,6 +79,11 @@ function restoreEnvironment(): void {
   process.env.RUNNER_TEMP = ORIGINAL_RUNNER_TEMP;
   process.env.UV_CACHE_DIR = ORIGINAL_UV_CACHE_DIR;
   process.env.UV_PYTHON = ORIGINAL_UV_PYTHON;
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
   process.env.UV_PYTHON_INSTALL_DIR = ORIGINAL_UV_PYTHON_INSTALL_DIR;
 }
 
@@ -98,6 +105,8 @@ describe("loadInputs", () => {
       source: CacheLocalSource.Default,
     });
     expect(inputs.pythonDir).toBe("/runner-temp/uv-python-dir");
+    expect(inputs.pythonArch).toBe("");
+    expect(inputs.exportPythonArch).toBe(false);
     expect(inputs.venvPath).toBe("/workspace/.venv");
     expect(inputs.manifestFile).toBeUndefined();
     expect(inputs.resolutionStrategy).toBe("highest");
@@ -149,6 +158,28 @@ describe("loadInputs", () => {
     expect(inputs.pythonVersion).toBe("");
   });
 
+  it("prefers the python-arch input over UV_PYTHON_ARCH", () => {
+    mockInputs["working-directory"] = "/workspace";
+    mockInputs["python-arch"] = "x86_64";
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    const inputs = loadInputs();
+    expect(inputs.pythonArch).toBe("x86_64");
+    expect(inputs.exportPythonArch).toBe(true);
+  });
+
+  it.each(["aarch64", "x86_64_v3", ""])(
+    "uses UV_PYTHON_ARCH when python-arch is omitted: %s",
+    (pythonArch) => {
+      mockInputs["working-directory"] = "/workspace";
+      process.env.UV_PYTHON_ARCH = pythonArch;
+
+      const inputs = loadInputs();
+      expect(inputs.pythonArch).toBe(pythonArch);
+      expect(inputs.exportPythonArch).toBe(false);
+    },
+  );
+
   it.each(["pull_request_target", "workflow_run", "release"])(
     "disables automatic caching for the %s event",
     (eventName) => {

__tests__/utils/python-arch.test.ts

@@ -0,0 +1,152 @@
+import { promisify } from "node:util";
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  jest,
+} from "@jest/globals";
+
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
+
+const mockExecFile =
+  jest.fn<
+    (...args: unknown[]) => Promise<{
+      stdout: string;
+      stderr: string;
+    }>
+  >();
+const mockExportVariable = jest.fn();
+
+jest.unstable_mockModule("node:child_process", () => ({
+  execFile: Object.assign(mockExecFile, { [promisify.custom]: mockExecFile }),
+}));
+jest.unstable_mockModule("@actions/core", () => ({
+  exportVariable: mockExportVariable,
+}));
+jest.unstable_mockModule("../../src/utils/logging", () => ({
+  info: jest.fn(),
+}));
+
+const { setupPythonArch } = await import("../../src/utils/python-arch");
+
+const supportedProbe = Object.assign(new Error("uv exited with code 2"), {
+  code: 2,
+  stderr:
+    "error: Failed to parse environment variable `UV_PYTHON_ARCH` with invalid value `setup-uv-probe`: Unknown architecture: setup-uv-probe\n",
+});
+const success = { stderr: "", stdout: "/cache\n" };
+
+beforeEach(() => {
+  delete process.env.UV_PYTHON_ARCH;
+  jest.clearAllMocks();
+  mockExecFile
+    .mockReset()
+    .mockRejectedValueOnce(supportedProbe)
+    .mockResolvedValueOnce(success);
+});
+
+afterEach(() => {
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
+});
+
+describe("setupPythonArch", () => {
+  it("leaves Python selection alone without an architecture", async () => {
+    await setupPythonArch("/tools/uv", "", false);
+
+    expect(mockExecFile).not.toHaveBeenCalled();
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("validates a resolved environment architecture without exporting it", async () => {
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    await setupPythonArch("/tools/uv", "aarch64", false);
+
+    expect(mockExecFile).toHaveBeenNthCalledWith(
+      2,
+      "/tools/uv",
+      ["--no-config", "cache", "dir"],
+      expect.objectContaining({
+        env: expect.objectContaining({ UV_PYTHON_ARCH: "aarch64" }),
+      }),
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+    expect(process.env.UV_PYTHON_ARCH).toBe("aarch64");
+  });
+
+  it.each(["/runner temp/uv", "C:\\runner temp\\uv.exe"])(
+    "validates and exports the architecture using the installed uv: %s",
+    async (uvPath) => {
+      await setupPythonArch(uvPath, "x86_64", true);
+
+      expect(mockExecFile).toHaveBeenNthCalledWith(
+        1,
+        uvPath,
+        ["--no-config", "cache", "dir"],
+        expect.objectContaining({
+          encoding: "utf8",
+          env: expect.objectContaining({ UV_PYTHON_ARCH: "setup-uv-probe" }),
+        }),
+      );
+      expect(mockExecFile).toHaveBeenNthCalledWith(
+        2,
+        uvPath,
+        ["--no-config", "cache", "dir"],
+        expect.objectContaining({
+          env: expect.objectContaining({ UV_PYTHON_ARCH: "x86_64" }),
+        }),
+      );
+      expect(mockExportVariable).toHaveBeenCalledWith(
+        "UV_PYTHON_ARCH",
+        "x86_64",
+      );
+    },
+  );
+
+  it("rejects uv versions that ignore UV_PYTHON_ARCH", async () => {
+    mockExecFile.mockReset().mockResolvedValue(success);
+
+    await expect(setupPythonArch("/tools/uv", "x86_64", true)).rejects.toThrow(
+      "The installed version of uv does not support UV_PYTHON_ARCH",
+    );
+    expect(mockExecFile).toHaveBeenCalledTimes(1);
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("reports unexpected probe failures", async () => {
+    mockExecFile.mockReset().mockRejectedValue(
+      Object.assign(new Error("uv exited with code 2"), {
+        code: 2,
+        stderr: "error: unrelated configuration error\n",
+      }),
+    );
+
+    await expect(setupPythonArch("/tools/uv", "x86_64", true)).rejects.toThrow(
+      "Failed to check uv's support for UV_PYTHON_ARCH: error: unrelated configuration error",
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("reports invalid architectures before exporting them", async () => {
+    mockExecFile
+      .mockReset()
+      .mockRejectedValueOnce(supportedProbe)
+      .mockRejectedValueOnce(
+        Object.assign(new Error("uv exited with code 2"), {
+          code: 2,
+          stderr: "error: Unknown architecture: invalid\n",
+        }),
+      );
+
+    await expect(setupPythonArch("/tools/uv", "invalid", true)).rejects.toThrow(
+      "Failed to set Python architecture to invalid: error: Unknown architecture: invalid",
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+});