Add a Python architecture input

__tests__/cache/restore-cache.test.ts

@@ -1,9 +1,20 @@
-import { beforeEach, describe, expect, it, jest } from "@jest/globals";
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  jest,
+} from "@jest/globals";
 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");
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
 
 jest.unstable_mockModule("@actions/cache", () => ({
   restoreCache: mockRestoreCache,
@@ -24,9 +35,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");
@@ -38,7 +49,19 @@ function cacheKeyOutput(): string {
 }
 
 beforeEach(() => {
+  delete process.env.UV_PYTHON_ARCH;
   jest.clearAllMocks();
+  mockGetArch.mockReturnValue("x86_64");
+  mockGetOSNameVersion.mockReturnValue("ubuntu-24.04");
+  mockGetPlatform.mockResolvedValue("unknown-linux-gnu");
+});
+
+afterEach(() => {
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
 });
 
 describe("restoreCache", () => {
@@ -71,4 +94,39 @@ describe("restoreCache", () => {
       "setup-uv-2-x86_64-unknown-linux-gnu-ubuntu-24.04-3.11-dependencyhash-tests-3.11",
     );
   });
+
+  it("includes an inherited Python architecture in the cache key", async () => {
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    await restoreCache(createSetupInputs(), "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`,
+      );
+    },
+  );
 });