diff --git a/uv/private/whl_install/repository.bzl b/uv/private/whl_install/repository.bzl index 244f5d1de..dec872513 100644 --- a/uv/private/whl_install/repository.bzl +++ b/uv/private/whl_install/repository.bzl @@ -10,7 +10,7 @@ produce a filegroup/TreeArtifact. load("@bazel_features//:features.bzl", features = "bazel_features") load("//uv/private:parse_whl_name.bzl", "parse_whl_name") load("//uv/private/constraints:defs.bzl", "MAJORS", "MINORS") -load("//uv/private/constraints/platform:defs.bzl", "supported_platform") +load("//uv/private/constraints/platform:defs.bzl", "MACOS_ARCH_GROUPS", "WINDOWS_PLATFORMS", "supported_platform") load("//uv/private/constraints/python:defs.bzl", "supported_python") load("//uv/private/pprint:defs.bzl", "pprint") @@ -446,6 +446,127 @@ def sort_select_arms(arms): pairs = sorted(arms.items(), key = lambda kv: select_key(kv[0]), reverse = True) return {a: b for a, b in pairs} +def _detect_host_os(os_name): + """Map repository_ctx.os.name to a wheel-platform family.""" + if os_name == "linux": + return "linux" + if os_name == "mac os x": + return "macos" + if os_name.startswith("windows"): + return "windows" + return None + +def _detect_host_arch(arch): + """Map repository_ctx.os.arch to the canonical wheel architecture suffix.""" + return { + "amd64": "x86_64", + "x86_64": "x86_64", + "aarch64": "aarch64", + "arm64": "aarch64", + }.get(arch) + +def _platform_tag_arch(platform_tag): + """Extract the canonical CPU architecture from a wheel platform tag. + + Returns None for tags that do not carry a parseable architecture (e.g. + bare "win32" or "any"). + """ + if platform_tag == "any": + return None + + # linux_x86_64, linux_aarch64, ... + if platform_tag.startswith("linux_"): + arch = platform_tag[len("linux_"):] + else: + # manylinux_2_17_x86_64, musllinux_1_2_aarch64, macosx_11_0_arm64, ... + parts = platform_tag.split("_", 3) + if len(parts) == 4: + arch = parts[3] + else: + return None + + # Normalize aliases so macOS "arm64" matches host "aarch64" and friends. + return { + "amd64": "x86_64", + "x86_64": "x86_64", + "aarch64": "aarch64", + "arm64": "aarch64", + }.get(arch, arch) + +def _platform_tag_matches_host(platform_tag, host_os, host_arch): + """Return True if a wheel platform tag is compatible with the host platform.""" + if platform_tag == "any": + return True + + if host_os == "linux": + if platform_tag.startswith("linux_"): + return _platform_tag_arch(platform_tag) == host_arch + if platform_tag.startswith("manylinux_") or platform_tag.startswith("musllinux_"): + return _platform_tag_arch(platform_tag) == host_arch + return False + + if host_os == "macos": + if not platform_tag.startswith("macosx_"): + return False + arch = _platform_tag_arch(platform_tag) + if arch == host_arch: + return True + group_arches = MACOS_ARCH_GROUPS.get(arch) + if group_arches: + return host_arch in group_arches + return False + + if host_os == "windows": + wheel_arch = WINDOWS_PLATFORMS.get(platform_tag) + return wheel_arch == host_arch + + return False + +def pick_gazelle_index_whl(repository_ctx, prebuilds, host_os = None, host_arch = None): + """Pick a wheel for Gazelle indexing without fetching foreign platforms. + + Any wheel of the package works for indexing, but prefer: + 1. A pure-Python ``any`` wheel — small and platform-independent. + 2. A wheel whose platform tags match the host OS/arch, because the build + or test is very likely to download that wheel already. + 3. The first wheel in the lockfile as a deterministic fallback. + + Args: + repository_ctx: the repository rule context (used for host detection when + ``host_os``/``host_arch`` are not supplied). + prebuilds: dict mapping wheel filename to wheel-file target label. + host_os: optional override for the host OS family ("linux", "macos", + "windows"); used by tests. + host_arch: optional override for the host CPU architecture ("x86_64", + "aarch64"); used by tests. + + Returns: + The selected target label string. + """ + if host_os == None: + host_os = _detect_host_os(repository_ctx.os.name) + if host_arch == None: + host_arch = _detect_host_arch(repository_ctx.os.arch) + + any_wheel = None + host_wheel = None + + for whl, target in prebuilds.items(): + parsed = parse_whl_name(whl) + for platform_tag in parsed.platform_tags: + if platform_tag == "any": + if any_wheel == None: + any_wheel = target + elif host_os and host_arch and _platform_tag_matches_host(platform_tag, host_os, host_arch): + if host_wheel == None: + host_wheel = target + break + if any_wheel and host_wheel: + # ``any`` is the cheapest possible index wheel; stop looking. + break + + return any_wheel or host_wheel or prebuilds.values()[0] + def compatible_python_tags(python_tag, abi_tag): if abi_tag != "abi3" or not python_tag.startswith("cp"): return [python_tag] @@ -613,7 +734,7 @@ py_library( ) if prebuilds: - gazelle_index_whl = prebuilds.values()[0] # Effectively random choice :shrug: + gazelle_index_whl = pick_gazelle_index_whl(repository_ctx, prebuilds) elif default_target: gazelle_index_whl = default_target else: diff --git a/uv/private/whl_install/test.bzl b/uv/private/whl_install/test.bzl index 7bb65a230..00be23999 100644 --- a/uv/private/whl_install/test.bzl +++ b/uv/private/whl_install/test.bzl @@ -3,7 +3,7 @@ load("@bazel_skylib//lib:unittest.bzl", "analysistest", "asserts", "unittest") load("@bazel_skylib//rules:write_file.bzl", "write_file") load("//py/private:providers.bzl", "PyWheelsInfo") -load(":repository.bzl", "compatible_python_tags", "parse_console_script", "parse_record_path", "select_key", "site_packages_segments", "sort_select_arms", "source_specificity") +load(":repository.bzl", "compatible_python_tags", "parse_console_script", "parse_record_path", "pick_gazelle_index_whl", "select_key", "site_packages_segments", "sort_select_arms", "source_specificity") load(":rule.bzl", "source_built_wheel", "whl_install") def _whl_sorting_test_impl(ctx): @@ -456,6 +456,145 @@ def metadata_selection_test_suite(name): leaked_top_levels = [], ) +# --- gazelle_index_whl host-aware selection --------------------------------- + +def _pick_gazelle_index_whl_prefers_any_test_impl(ctx): + env = unittest.begin(ctx) + + prebuilds = { + "pkg-1.0-cp311-cp311-manylinux_2_17_x86_64.whl": "@linux//file", + "pkg-1.0-py3-none-any.whl": "@any//file", + "pkg-1.0-cp311-cp311-macosx_11_0_arm64.whl": "@mac//file", + } + asserts.equals( + env, + "@any//file", + pick_gazelle_index_whl(None, prebuilds, host_os = "linux", host_arch = "x86_64"), + ) + + return unittest.end(env) + +pick_gazelle_index_whl_prefers_any_test = unittest.make( + _pick_gazelle_index_whl_prefers_any_test_impl, +) + +def _pick_gazelle_index_whl_linux_test_impl(ctx): + env = unittest.begin(ctx) + + # Prefers the host-arch manylinux wheel over a foreign macOS wheel. + prebuilds_manylinux = { + "pkg-1.0-cp311-cp311-macosx_11_0_arm64.whl": "@mac//file", + "pkg-1.0-cp311-cp311-manylinux_2_17_x86_64.whl": "@linux//file", + } + asserts.equals( + env, + "@linux//file", + pick_gazelle_index_whl(None, prebuilds_manylinux, host_os = "linux", host_arch = "x86_64"), + ) + + # musllinux also matches a Linux host. + prebuilds_musl = { + "pkg-1.0-cp311-cp311-musllinux_1_2_aarch64.whl": "@musl//file", + "pkg-1.0-cp311-cp311-manylinux_2_17_x86_64.whl": "@linux//file", + } + asserts.equals( + env, + "@musl//file", + pick_gazelle_index_whl(None, prebuilds_musl, host_os = "linux", host_arch = "aarch64"), + ) + + # bare linux_* tag. + prebuilds_bare = { + "pkg-1.0-cp311-cp311-linux_aarch64.whl": "@bare//file", + } + asserts.equals( + env, + "@bare//file", + pick_gazelle_index_whl(None, prebuilds_bare, host_os = "linux", host_arch = "aarch64"), + ) + + return unittest.end(env) + +pick_gazelle_index_whl_linux_test = unittest.make( + _pick_gazelle_index_whl_linux_test_impl, +) + +def _pick_gazelle_index_whl_macos_test_impl(ctx): + env = unittest.begin(ctx) + + # arm64 host prefers the arm64-specific wheel. + prebuilds_arm = { + "pkg-1.0-cp311-cp311-macosx_10_9_x86_64.whl": "@mac_x86//file", + "pkg-1.0-cp311-cp311-macosx_11_0_arm64.whl": "@mac_arm//file", + } + asserts.equals( + env, + "@mac_arm//file", + pick_gazelle_index_whl(None, prebuilds_arm, host_os = "macos", host_arch = "aarch64"), + ) + + # universal2 covers both architectures. + prebuilds_universal = { + "pkg-1.0-cp311-cp311-macosx_11_0_universal2.whl": "@mac_univ//file", + "pkg-1.0-cp311-cp311-macosx_11_0_arm64.whl": "@mac_arm//file", + } + asserts.equals( + env, + "@mac_univ//file", + pick_gazelle_index_whl(None, prebuilds_universal, host_os = "macos", host_arch = "x86_64"), + ) + + return unittest.end(env) + +pick_gazelle_index_whl_macos_test = unittest.make( + _pick_gazelle_index_whl_macos_test_impl, +) + +def _pick_gazelle_index_whl_windows_test_impl(ctx): + env = unittest.begin(ctx) + + prebuilds = { + "pkg-1.0-cp311-cp311-win_amd64.whl": "@win_amd64//file", + "pkg-1.0-cp311-cp311-win_arm64.whl": "@win_arm64//file", + "pkg-1.0-cp311-cp311-win32.whl": "@win32//file", + } + asserts.equals( + env, + "@win_amd64//file", + pick_gazelle_index_whl(None, prebuilds, host_os = "windows", host_arch = "x86_64"), + ) + asserts.equals( + env, + "@win_arm64//file", + pick_gazelle_index_whl(None, prebuilds, host_os = "windows", host_arch = "aarch64"), + ) + + return unittest.end(env) + +pick_gazelle_index_whl_windows_test = unittest.make( + _pick_gazelle_index_whl_windows_test_impl, +) + +def _pick_gazelle_index_whl_fallback_test_impl(ctx): + env = unittest.begin(ctx) + + # No host match: falls back to the first wheel in the lockfile order. + prebuilds = { + "pkg-1.0-cp311-cp311-macosx_11_0_arm64.whl": "@mac//file", + "pkg-1.0-cp311-cp311-win_amd64.whl": "@win//file", + } + asserts.equals( + env, + "@mac//file", + pick_gazelle_index_whl(None, prebuilds, host_os = "linux", host_arch = "x86_64"), + ) + + return unittest.end(env) + +pick_gazelle_index_whl_fallback_test = unittest.make( + _pick_gazelle_index_whl_fallback_test_impl, +) + def whl_install_suite(): unittest.suite( "whl_sorting_tests", @@ -481,3 +620,11 @@ def whl_install_suite(): "console_script_tests", console_script_test, ) + unittest.suite( + "pick_gazelle_index_whl_tests", + pick_gazelle_index_whl_prefers_any_test, + pick_gazelle_index_whl_linux_test, + pick_gazelle_index_whl_macos_test, + pick_gazelle_index_whl_windows_test, + pick_gazelle_index_whl_fallback_test, + )