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__/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();
+  });
+});