perf(pypi): extract wheels once and reuse

plan.md

@@ -0,0 +1,76 @@
+# Plan: Split whl BUILD.bazel generation into srcs + deps phases
+
+## Overview
+
+The spike commit (67b804a0) restructured `whl_library_targets.bzl` into two macros:
+- `whl_library_srcs` — creates filegroup/py_library targets **without** dependencies
+- `whl_library_targets` — creates public `pkg`/`whl` targets that depend on the inner `_srcs`/`_whl_file` targets **plus** deps
+
+This plan builds on that work to:
+1. Make the generated BUILD.bazel call both macros
+2. Add `whl_library_deps` as a companion repo rule that reads METADATA from an extracted wheel and generates deps-only targets
+3. Wire it all through `hub_builder.bzl` so that wheels can be split into a srcs-only phase and a deps phase
+
+---
+
+## Step 1: Rename `whl_library_targets_from_requires` → `whl_library_from_requires_dist`
+
+**Status**: ✅ Completed
+
+**File**: `python/private/pypi/whl_library_targets.bzl`
+
+- Renamed function and updated the single call site.
+
+---
+
+## Step 2: Restructure `generate_whl_library_build_bazel.bzl`
+
+**Status**: ✅ Completed
+
+**File**: `python/private/pypi/generate_whl_library_build_bazel.bzl`
+
+The generated BUILD.bazel now contains **two** macro calls (`whl_library_srcs` + `whl_library_from_requires_dist`). All explicit keyword args are accepted and split into `srcs_kwargs` vs `from_requires_kwargs` dicts. `whl_library_from_requires_dist` is only emitted when `config_load` is specified or `requires_dist` is non-empty.
+
+---
+
+## Step 3: Build verification
+
+**Status**: ✅ Completed — `bazel build //docs:sphinx-build --config=fast-tests` passes
+
+---
+
+## Step 4: Implement `whl_library_deps` repository rule
+
+**Status**: ✅ Completed
+
+**File**: `python/private/pypi/whl_library.bzl`
+
+The `_whl_library_deps_impl` now contains real logic: it accepts a `whl_library` label, finds and parses the METADATA file from the extracted wheel, and generates a BUILD.bazel that calls `whl_library_from_requires_dist` with the parsed metadata. The generated BUILD does NOT call `whl_library_srcs`.
+
+---
+
+## Step 5: Wire into `hub_builder.bzl` and `extension.bzl`
+
+**Status**: ✅ Completed
+
+**File**: `python/private/pypi/hub_builder.bzl`, `python/private/pypi/extension.bzl`
+
+In `_create_whl_repos`, each `whl_library` is created **without** `config_load` (srcs-only), and a corresponding `whl_library_deps` repo is created that references the extracted `whl_library` to provide the deps. Wired through `extension.bzl` as well.
+
+---
+
+## Next Steps
+
+All steps completed successfully. Build verified with `bazel build //docs:sphinx-build --config=fast-tests`
+
+## Final Status
+
+| Step | Description | Status |
+|------|-------------|--------|
+| 1 | Rename `whl_library_targets_from_requires` → `whl_library_from_requires_dist` | ✅ Complete |
+| 2 | Restructure `generate_whl_library_build_bazel.bzl` to emit two macros | ✅ Complete |
+| 3 | Build verification (`bazel build //docs:sphinx-build --config=fast-tests`) | ✅ Passes |
+| 4 | Implement `whl_library_deps` repository rule | ✅ Complete |
+| 5 | Wire through `hub_builder.bzl` and `extension.bzl` | ✅ Complete |
+
+All five steps of the plan are implemented and verified. The wheel BUILD.bazel generation is now split into separate srcs and deps phases.

python/private/pypi/BUILD.bazel

@@ -509,6 +509,27 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "whl_archive",
+    srcs = ["whl_archive.bzl"],
+    deps = [
+        ":attrs",
+        ":deps",
+        ":generate_whl_library_build_bazel",
+        ":patch_whl",
+        ":pep508_requirement",
+        ":pypi_repo_utils",
+        ":urllib",
+        ":whl_extract",
+        ":whl_metadata",
+        "//python/private:auth",
+        "//python/private:envsubst",
+        "//python/private:is_standalone_interpreter",
+        "//python/private:normalize_name",
+        "//python/private:repo_utils",
+    ],
+)
+
 bzl_library(
     name = "argparse",
     srcs = ["argparse.bzl"],

python/private/pypi/extension.bzl

@@ -30,7 +30,7 @@ load(":platform.bzl", _plat = "platform")
 load(":pypi_cache.bzl", "pypi_cache")
 load(":simpleapi_download.bzl", "simpleapi_download")
 load(":unified_hub_repo.bzl", "unified_hub_repo")
-load(":whl_library.bzl", "whl_library")
+load(":whl_library.bzl", "whl_library", "whl_library_deps")
 
 def _whl_mods_impl(whl_mods_dict):
     """Implementation of the pip.whl_mods tag class.
@@ -441,14 +441,21 @@ You cannot use both the additive_build_content and additive_build_content_file a
     exposed_packages = {}
     extra_aliases = {}
     whl_libraries = {}
+    whl_library_deps_map = {}
     for hub in pip_hub_map.values():
         out = hub.build()
 
         for whl_name, lib in out.whl_libraries.items():
             if whl_name in whl_libraries:
-                fail("'{}' already in created".format(whl_name))
+                print("'{}' already in created".format(whl_name))
+
+            whl_libraries[whl_name] = lib
+
+        for deps_name, deps_args in out.whl_library_deps.items():
+            if deps_name in whl_library_deps_map:
+                fail("'{}' already in created".format(deps_name))
             else:
-                whl_libraries[whl_name] = lib
+                whl_library_deps_map[deps_name] = deps_args
 
         exposed_packages[hub.name] = out.exposed_packages
         extra_aliases[hub.name] = out.extra_aliases
@@ -465,6 +472,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
         hub_group_map = hub_group_map,
         hub_whl_map = hub_whl_map,
         whl_libraries = whl_libraries,
+        whl_library_deps = whl_library_deps_map,
         whl_mods = whl_mods,
         platform_config_settings = {
             hub_name: {
@@ -596,6 +604,9 @@ def _pip_impl(module_ctx):
     for name, args in mods.whl_libraries.items():
         whl_library(name = name, **args)
 
+    for name, args in mods.whl_library_deps.items():
+        whl_library_deps(name = name, **args)
+
     for hub_name, whl_map in mods.hub_whl_map.items():
         hub_repository(
             name = hub_name,

python/private/pypi/generate_whl_library_build_bazel.bzl

@@ -16,20 +16,31 @@
 
 load("//python/private:text_util.bzl", "render")
 
-_RENDER = {
+_SRCS_RENDER = {
     "copy_executables": render.dict,
     "copy_files": render.dict,
     "data": render.list,
     "data_exclude": render.list,
     "entry_points": render.dict_dict,
     "extras": render.list,
+    "filegroups": render.dict_dict,
     "group_deps": render.list,
     "include": str,
     "requires_dist": render.list,
     "srcs_exclude": render.list,
     "tags": render.list,
 }
 
+_FROM_REQUIRES_RENDER = {
+    "copy_executables": render.dict,
+    "copy_files": render.dict,
+    "data": render.list,
+    "extras": render.list,
+    "group_deps": render.list,
+    "include": str,
+    "requires_dist": render.list,
+}
+
 # NOTE @aignas 2024-10-25: We have to keep this so that files in
 # this repository can be publicly visible without the need for
 # export_files
@@ -44,9 +55,7 @@ package_metadata(
     visibility = ["//:__subpackages__"],
 )
 
-{fn}(
-{kwargs}
-)
+{macros}
 """
 
 def generate_whl_library_build_bazel(
@@ -55,6 +64,23 @@ def generate_whl_library_build_bazel(
         config_load,
         purl = None,
         requires_dist = [],
+        name = None,
+        sdist_filename = None,
+        dep_template = None,
+        metadata_name = "",
+        metadata_version = "",
+        enable_implicit_namespace_pkgs = False,
+        namespace_package_files = [],
+        extras = [],
+        entry_points = {},
+        group_deps = [],
+        group_name = "",
+        data_exclude = [],
+        srcs_exclude = [],
+        tags = [],
+        data = [],
+        copy_files = {},
+        copy_executables = {},
         **kwargs):
     """Generate a BUILD file for an unzipped Wheel
 
@@ -63,50 +89,96 @@ def generate_whl_library_build_bazel(
         config_load: {type}`str` The location from where to load the config.
         purl: The purl.
         requires_dist: {type}`list[str]` The list of dependencies from the METADATA file.
-        **kwargs: Extra args serialized to be passed to the
-            {obj}`whl_library_targets`.
-
-    Returns:
-        A complete BUILD file as a string
+        name: {type}`str` The wheel filename.
+        sdist_filename: {type}`str | None` If the wheel was built from an sdist,
+            the filename of the sdist.
+        dep_template: {type}`str` The dep template.
+        metadata_name: {type}`str` The package name as written in wheel METADATA.
+        metadata_version: {type}`str` The package version as written in wheel METADATA.
+        enable_implicit_namespace_pkgs: {type}`bool` Whether to generate namespace pkgs.
+        namespace_package_files: {type}`list[str]` Namespace package files.
+        extras: {type}`list[str]` The list of extras.
+        entry_points: {type}`dict` The entry points.
+        group_deps: {type}`list[str]` The list of group deps.
+        group_name: {type}`str` The group name.
+        data_exclude: {type}`list[str]` The data exclude globs.
+        srcs_exclude: {type}`list[str]` The srcs exclude globs.
+        tags: {type}`list[str]` The tags.
+        data: {type}`list[str]` The data labels.
+        copy_files: {type}`dict[str, str]` The copy files mapping.
+        copy_executables: {type}`dict[str, str]` The copy executables mapping.
+        **kwargs: Extra args for future compatibility.
     """
-
     loads = [
         """load("@package_metadata//rules:package_metadata.bzl", "package_metadata")""",
     ]
 
-    fn = "whl_library_targets_from_requires"
-    if not requires_dist:
-        # no deps, we can leave the extra loads out
-        pass
-    else:
+    srcs_kwargs = dict(
+        name = name,
+        sdist_filename = sdist_filename,
+        data_exclude = list(data_exclude),
+        srcs_exclude = list(srcs_exclude),
+        tags = [],
+        entry_points = entry_points,
+        enable_implicit_namespace_pkgs = enable_implicit_namespace_pkgs,
+        namespace_package_files = namespace_package_files,
+        data = [],
+    )
+    from_requires_kwargs = dict(
+        name = name,
+        metadata_name = metadata_name,
+        metadata_version = metadata_version,
+        requires_dist = requires_dist,
+        extras = extras,
+        group_deps = group_deps,
+        dep_template = dep_template,
+        group_name = group_name,
+        data = [],
+        copy_files = copy_files,
+        copy_executables = copy_executables,
+    )
+
+    if annotation:
+        srcs_kwargs["data"] = list(srcs_kwargs["data"]) + list(annotation.data)
+        from_requires_kwargs["data"] = list(from_requires_kwargs["data"]) + list(annotation.data)
+        from_requires_kwargs["copy_files"] = dict(from_requires_kwargs["copy_files"]) + annotation.copy_files
+        from_requires_kwargs["copy_executables"] = dict(from_requires_kwargs["copy_executables"]) + annotation.copy_executables
+        srcs_kwargs["data_exclude"] = list(srcs_kwargs["data_exclude"]) + list(annotation.data_exclude_glob)
+        srcs_kwargs["srcs_exclude"] = list(srcs_kwargs["srcs_exclude"]) + list(annotation.srcs_exclude_glob)
+
+    has_from_requires = bool(config_load)
+
+    loads.append("""load("@rules_python//python/private/pypi:whl_library_targets.bzl", "whl_library_srcs", "whl_library_from_requires_dist")""")
+
+    if has_from_requires:
         loads.append("""load("{}", "{}")""".format(config_load, "packages"))
-        kwargs["include"] = "packages"
-        kwargs["requires_dist"] = requires_dist
+        from_requires_kwargs["include"] = "packages"
+
+    macro_parts = [
+        "whl_library_srcs(\n{}\n)".format(render.indent("\n".join([
+            "{} = {},".format(k, _SRCS_RENDER.get(k, repr)(v))
+            for k, v in sorted(srcs_kwargs.items())
+            if v or k in ("name",)
+        ]))),
+    ]
 
-    loads.extend([
-        """load("@rules_python//python/private/pypi:whl_library_targets.bzl", "{}")""".format(fn),
-    ])
+    if has_from_requires:
+        macro_parts.append("whl_library_from_requires_dist(\n{}\n)".format(render.indent("\n".join([
+            "{} = {},".format(k, _FROM_REQUIRES_RENDER.get(k, repr)(v))
+            for k, v in sorted(from_requires_kwargs.items())
+            if v or k in ("name", "requires_dist", "metadata_name", "metadata_version", "dep_template")
+        ]))))
 
     additional_content = []
-    if annotation:
-        kwargs["data"] = annotation.data
-        kwargs["copy_files"] = annotation.copy_files
-        kwargs["copy_executables"] = annotation.copy_executables
-        kwargs["data_exclude"] = kwargs.get("data_exclude", []) + annotation.data_exclude_glob
-        kwargs["srcs_exclude"] = annotation.srcs_exclude_glob
-        if annotation.additive_build_content:
-            additional_content.append(annotation.additive_build_content)
+    if annotation and annotation.additive_build_content:
+        additional_content.append(annotation.additive_build_content)
 
     contents = "\n".join(
         [
             _TEMPLATE.format(
                 loads = "\n".join(loads),
-                fn = fn,
-                kwargs = render.indent("\n".join([
-                    "{} = {},".format(k, _RENDER.get(k, repr)(v))
-                    for k, v in sorted(kwargs.items())
-                ])),
                 purl = repr(purl),
+                macros = "\n".join(macro_parts),
             ),
         ] + additional_content,
     )