diff --git a/uv/private/extension/defs.bzl b/uv/private/extension/defs.bzl index 1379fce00..cbc493fa8 100644 --- a/uv/private/extension/defs.bzl +++ b/uv/private/extension/defs.bzl @@ -87,6 +87,29 @@ def _deduplicate_whl_files(whls): whl_files.append(whl) return whl_files +def group_whls_by_target(whls): + """Group every wheel filename that resolves to the same fetched label.""" + grouped = {} + for whl_name, whl_target in whls.items(): + if whl_target: + grouped.setdefault(whl_target, {})[whl_name] = whl_target + return grouped + +def wheel_variant_suffix(whl_target): + """Return the stable hash suffix from a generated wheel repository label.""" + return whl_target.split("__")[-1].split("//")[0] + +def shared_whl_variant_key(install_cfg, variant_whls): + """Return an order-sensitive identity for an unpatched wheel variant.""" + if install_cfg.post_install_patches: + return None + return json.encode([ + install_cfg.metadata_directory, + # Wheel selection keeps the first entry when specificity ties. + variant_whls.items(), + getattr(install_cfg, "exclude_glob", []), + ]) + def parse_declared_console_script(name, entry_point): """Canonicalize one override_package console-script declaration. @@ -739,16 +762,48 @@ def _uv_impl(module_ctx): sbuild_kwargs["resource_set"] = sbuild_cfg.resource_set sdist_build(**sbuild_kwargs) + shared_variant_installs = {} for install_id, install_cfg in cfg.install_cfgs.items(): + # Metadata extraction belongs in one repository per fetched wheel. + # The public install repository can then select one of these variants + # without fetching metadata for every inactive platform. + variant_installs = {} + for whl_target, variant_whls in group_whls_by_target(install_cfg.whls).items(): + shared_key = shared_whl_variant_key(install_cfg, variant_whls) + if shared_key in shared_variant_installs: + variant_installs[whl_target] = shared_variant_installs[shared_key] + continue + variant_name = "{}__variant_{}".format( + install_id, + wheel_variant_suffix(whl_target), + ) + variant_kwargs = { + "metadata_directory": install_cfg.metadata_directory, + "name": variant_name, + "sbuild": None, + "sbuild_console_scripts": [], + "whls": json.encode(variant_whls), + "whl_files": [whl_target], + } + if install_cfg.post_install_patches: + variant_kwargs["post_install_patches"] = json.encode(install_cfg.post_install_patches) + variant_kwargs["post_install_patch_strip"] = install_cfg.post_install_patch_strip + exclude_glob = getattr(install_cfg, "exclude_glob", []) + if exclude_glob: + variant_kwargs["exclude_glob"] = exclude_glob + whl_install(**variant_kwargs) + variant_installs[whl_target] = "@{}//:install".format(variant_name) + if shared_key: + shared_variant_installs[shared_key] = variant_installs[whl_target] install_kwargs = { "metadata_directory": install_cfg.metadata_directory, "name": install_id, "sbuild": install_cfg.sbuild, + "variant_installs": json.encode(variant_installs), "sbuild_console_scripts": install_cfg.sbuild_console_scripts, "whls": json.encode(install_cfg.whls), - # Parallel list of the same wheel labels as a real label_list, - # so the whl_install repo rule can `rctx.path()` them to peek - # at `*.dist-info/RECORD` for top-level metadata. + # Keep this for callers that instantiate whl_install directly + # without variant repos. "whl_files": _deduplicate_whl_files(install_cfg.whls.values()), } if install_cfg.post_install_patches: diff --git a/uv/private/extension/test_defs.bzl b/uv/private/extension/test_defs.bzl index 05ae3bff3..58bd6918c 100644 --- a/uv/private/extension/test_defs.bzl +++ b/uv/private/extension/test_defs.bzl @@ -1,7 +1,7 @@ """Unit tests for helpers in defs.bzl""" load("@bazel_skylib//lib:unittest.bzl", "asserts", "unittest") -load(":defs.bzl", "parse_declared_console_script") +load(":defs.bzl", "parse_declared_console_script", "shared_whl_variant_key") load(":lockfile.bzl", "url_basename") def _url_basename_test_impl(ctx): @@ -72,6 +72,49 @@ def _declared_console_script_test_impl(ctx): declared_console_script_test = unittest.make(_declared_console_script_test_impl) +def _shared_whl_variant_key_test_impl(ctx): + env = unittest.begin(ctx) + whl_target = "@whl__demo__abcdef//file" + ordered = { + "demo-1.0-1-py3-none-any.whl": whl_target, + "demo-1.0-2-py3-none-any.whl": whl_target, + } + reversed_order = { + "demo-1.0-2-py3-none-any.whl": whl_target, + "demo-1.0-1-py3-none-any.whl": whl_target, + } + config = struct( + exclude_glob = [], + metadata_directory = "demo-1.0.dist-info", + post_install_patches = [], + ) + + key = shared_whl_variant_key(config, ordered) + asserts.equals(env, key, shared_whl_variant_key(config, ordered)) + asserts.equals(env, False, key == shared_whl_variant_key(config, reversed_order)) + asserts.equals( + env, + False, + key == shared_whl_variant_key(struct( + exclude_glob = ["**/*.pyi"], + metadata_directory = "demo-1.0.dist-info", + post_install_patches = [], + ), ordered), + ) + asserts.equals( + env, + None, + shared_whl_variant_key(struct( + exclude_glob = [], + metadata_directory = "demo-1.0.dist-info", + post_install_patches = ["//:demo.patch"], + ), ordered), + ) + + return unittest.end(env) + +shared_whl_variant_key_test = unittest.make(_shared_whl_variant_key_test_impl) + def defs_test_suite(): unittest.suite( "url_basename_tests", @@ -81,3 +124,7 @@ def defs_test_suite(): "declared_console_script_tests", declared_console_script_test, ) + unittest.suite( + "shared_whl_variant_key_tests", + shared_whl_variant_key_test, + ) diff --git a/uv/private/whl_install/repository.bzl b/uv/private/whl_install/repository.bzl index 91e75f201..13213c913 100644 --- a/uv/private/whl_install/repository.bzl +++ b/uv/private/whl_install/repository.bzl @@ -685,6 +685,95 @@ filegroup( "//conditions:default": "checked-hash", })""" + # Keep the public package repository as a cheap selector, while moving + # metadata extraction into one repository per concrete fetched wheel. + # This lets Bazel fetch only the selected variant's repository. + variant_installs = json.decode(repository_ctx.attr.variant_installs) if repository_ctx.attr.variant_installs else {} + if variant_installs: + install_arms = { + condition: variant_installs[target] + for condition, target in select_arms.items() + } + default_install = ":whl_missing" if prebuilds else None + if sbuild_target: + fallback_install_attrs = """ + src = ":source_built_wheel", + compile_pyc = {compile_pyc}, + pyc_invalidation_mode = {pyc_invalidation_mode},""".format( + compile_pyc = compile_pyc_select, + pyc_invalidation_mode = pyc_invalidation_mode_select, + ) + if post_install_patches: + fallback_install_attrs += """ + patches = {patches}, + patch_strip = {strip},""".format( + patches = repr(post_install_patches), + strip = post_install_patch_strip, + ) + exclude_glob = getattr(repository_ctx.attr, "exclude_glob", []) + if exclude_glob: + fallback_install_attrs += """ + exclude_glob = {exclude_glob},""".format( + exclude_glob = indent(pprint(exclude_glob), " " * 4).lstrip(), + ) + content.append( + """ +whl_install( + name = "actual_install",{attrs} + visibility = ["//visibility:private"], +)""".format(attrs = fallback_install_attrs), + ) + default_install = ":actual_install" + + content.append( + """ +select_chain( + name = "selected_install", + arms = {arms}, + default_target = {default_target}, + visibility = ["//visibility:private"], +) +""".format( + arms = indent(pprint(install_arms), " ").lstrip(), + default_target = repr(default_install), + ), + ) + + if extra_deps or extra_data: + content.append( + """ +py_library( + name = "install", + srcs = [], + deps = [":selected_install"] + {extra_deps}, + data = {extra_data}, + visibility = ["//visibility:public"], +) +""".format( + extra_deps = repr(extra_deps), + extra_data = repr(extra_data), + ), + ) + else: + content.append( + """ +alias( + name = "install", + actual = ":selected_install", + visibility = ["//visibility:public"], +) +""", + ) + + content.append(""" +exports_files( + ["BUILD.bazel"], + visibility = ["//visibility:public"], +) +""") + repository_ctx.file("BUILD.bazel", content = "\n".join(content)) + return repository_ctx.repo_metadata(reproducible = True) + # Peek into each selectable wheel to extract the top-level names it # installs AND its `[console_scripts]` entry points, keyed by the # wheel's file basename. This powers PyWheelsInfo, which py_binary @@ -905,6 +994,7 @@ whl_install = repository_rule( # `-.dist-info`, the directory `_extract_wheel_metadata` # strips to when extracting RECORD/entry_points.txt out of the wheel. "metadata_directory": attr.string(), + "variant_installs": attr.string(default = ""), "whls": attr.string(), # Mirror of the http_file labels from `whls`, declared as a real # label_list so Bazel adds those repos to this repo's visibility