wip

aignas · aignas · commit 3eaf55d43dab · 2026-07-04T13:12:59.000+09:00
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
@@ -31,7 +31,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.
@@ -459,14 +459,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
@@ -483,6 +490,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: {
@@ -614,6 +622,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,
diff --git a/python/private/pypi/generate_whl_library_build_bazel.bzl b/python/private/pypi/generate_whl_library_build_bazel.bzl
@@ -16,7 +16,9 @@
 
 load("//python/private:text_util.bzl", "render")
 
-_RENDER = {
+# These are functions on how to render particular args, should be reused across all rendering
+# invocations to make things easier.
+_RENDER_FNS = {
     "copy_executables": render.dict,
     "copy_files": render.dict,
     "data": render.list,
@@ -29,6 +31,11 @@ _RENDER = {
     "srcs_exclude": render.list,
     "tags": render.list,
 }
+def _render(**kwargs):
+    return {
+        arg: _RENDER_FNS.get(arg, repr)(value)
+        for arg, value in kwargs.items
+    }
 
 # NOTE @aignas 2024-10-25: We have to keep this so that files in
 # this repository can be publicly visible without the need for
@@ -38,19 +45,12 @@ _TEMPLATE = """\
 
 package(default_visibility = ["//visibility:public"])
 
-package_metadata(
-    name = "package_metadata",
-    purl = {purl},
-    visibility = ["//:__subpackages__"],
-)
-
-{fn}(
-{kwargs}
-)
+{macros}
 """
 
 def generate_whl_library_build_bazel(
         *,
+        # TODO @aignas 2026-07-04: add extra args that are used in this function
         annotation = None,
         config_load,
         purl = None,
@@ -72,21 +72,41 @@ def generate_whl_library_build_bazel(
 
     loads = [
         """load("@package_metadata//rules:package_metadata.bzl", "package_metadata")""",
+        """load("@rules_python//python/private/pypi:whl_library_targets.bzl", "whl_library_srcs", "whl_library_from_requires_dist")"""
     ]
 
-    fn = "whl_library_targets_from_requires"
-    if not requires_dist:
-        # no deps, we can leave the extra loads out
-        pass
-    else:
-        loads.append("""load("{}", "{}")""".format(config_load, "packages"))
-        kwargs["include"] = "packages"
-        kwargs["requires_dist"] = requires_dist
-
-    loads.extend([
-        """load("@rules_python//python/private/pypi:whl_library_targets.bzl", "{}")""".format(fn),
-    ])
+    srcs_kwargs = dict(
+        name = name,
+        data = [],
+        sdist_filename = sdist_filename,
+        data_exclude = list(data_exclude),
+        srcs_exclude = list(srcs_exclude),
+        tags = [
+            "pypi_name={}".format(metadata_name),
+            "pypi_version={}".format(metadata_version),
+        ],
+        entry_points = entry_points,
+        enable_implicit_namespace_pkgs = enable_implicit_namespace_pkgs,
+        copy_files = copy_files,
+        copy_executables = copy_executables,
+        namespace_package_files = namespace_package_files,
+        data = [],
+        visibility = visibility,
+    )
+    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,
+    )
 
+    # NOTE, if users specify annotations, the wheel downloads are not reused this
+    # is to ensure that we don't break users config and also to ensure that we
+    # can have predictable results.
     additional_content = []
     if annotation:
         kwargs["data"] = annotation.data
@@ -97,19 +117,96 @@ def generate_whl_library_build_bazel(
         if annotation.additive_build_content:
             additional_content.append(annotation.additive_build_content)
 
+    macro_parts = [
+        render.call(
+            "package_metadata",
+            **_render(
+                name = "package_metadata",
+                purl = purl,
+                visibility = ["//:__subpackages__"],
+            ),
+        ),
+        render.call(
+            "whl_library_srcs",
+            **_render(**srcs_kwargs)
+        )
+    ]
+
+    if config_load:
+        loads.append("""load("{}", "{}")""".format(config_load, "packages"))
+        from_requires_kwargs["include"] = "packages"
+
+    macro_parts.append(render.call(
+        "whl_library_from_requires_dist",
+        **_render(**from_requires_kwargs),
+    ))
+
     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\n".join(macro_parts),
             ),
         ] + additional_content,
     )
 
     # NOTE: Ensure that we terminate with a new line
     return contents.rstrip() + "\n"
+
+def generate_whl_library_deps_build_bazel(
+        *,
+        # TODO @aignas 2026-07-04: add extra args that are used in this function
+        **kwargs):
+    """Generate a BUILD file for an unzipped Wheel
+
+
+    Returns:
+        A complete BUILD file as a string
+    """
+
+    loads = [
+        """load("@rules_python//python/private/pypi:whl_library_targets.bzl", "whl_library_from_requires_dist")"""
+    ]
+
+    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,
+    )
+
+    macro_parts = [
+        render.call(
+            "alias",
+            name=target,
+            actual=whl_library.same_package_label(target),
+        )
+        for target in [
+            # TODO @aignas 2026-07-04: use ./labels.bzl for the following
+            "package_metadata",
+            "data",
+            "dist_info",
+            "extracted_whl_files",
+        ]
+    ]
+
+    if config_load:
+        loads.append("""load("{}", "{}")""".format(config_load, "packages"))
+        from_requires_kwargs["include"] = "packages"
+
+    macro_parts.append(render.call(
+        "whl_library_from_requires_dist",
+        **_render(**from_requires_kwargs),
+    ))
+
+    contents = _TEMPLATE.format(
+        loads = "\n".join(loads),
+        macros = "\n\n".join(macro_parts),
+    )
+
+    # NOTE: Ensure that we terminate with a new line
+    return contents.rstrip() + "\n"
diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
@@ -71,6 +71,9 @@ def hub_builder(
         # Mapping of whl_library repo names and their kwargs.
         # dict[str repo_name, dict[str, object] kwargs]
         _whl_libraries = {},  # modified by _add_whl_library
+        # Mapping of whl_library_deps repo names and their kwargs.
+        # dict[str repo_name, dict[str, object] kwargs]
+        _whl_library_deps = {},  # modified by _add_whl_library
         # Map of repos and their config settings, and repo the config
         # setting originated from.
         # dict[str whl_name, dict[str config_setting, str repo_name]]
@@ -113,6 +116,7 @@ def _build(self):
         extra_aliases = {},
         exposed_packages = [],
         whl_libraries = {},
+        whl_library_deps = {},
     )
     if self._logger.failed():
         return ret
@@ -142,6 +146,10 @@ def _build(self):
         # Mapping of whl_library repo names and their kwargs.
         # dict[str repo_name, dict[str, object] kwargs]
         whl_libraries = self._whl_libraries,
+
+        # Mapping of whl_library_deps repo names and their kwargs.
+        # dict[str repo_name, dict[str, object] kwargs]
+        whl_library_deps = self._whl_library_deps,
     )
 
 def _pip_parse(self, module_ctx, pip_attr, python_version = None):
@@ -308,17 +316,43 @@ def _add_whl_library(self, *, python_version, whl, repo):
         # are more platforms defined than there are wheels for and users
         # disallow building from sdist.
         return
+    
+    forbidden_args = {
+        "annotation": True,
+        "extra_pip_args": True,
+        "python_interpreter": True,
+        "python_interpreter_target": True,
+    }
 
-    # TODO @aignas 2025-06-29: we should not need the version in the repo_name if
-    # we are using pipstar and we are downloading the wheel using the downloader
-    #
-    # However, for that we should first have a different way to reference closures with
-    # extras. For example, if some package depends on `foo[extra]` and another depends on
-    # `foo`, we should have 2 py_library targets.
-    repo_name = "{}_{}_{}".format(self.name, version_label(python_version), repo.repo_name)
+    if [arg for arg in repo.args if arg in forbidden_args]:
+        # no reuse of the whl_library because we have args that force the extraction of the whl
+        # in the hub context. If we have whl-only pipstar extraction, then we can reuse the
+        # extracted sources.
+        repos_dict = self._whl_libraries
+        deps_args = repo.args
+    else:
+        whl_repo_name = "whl_{}".format(repo.whl_repo_name)
 
-    if repo_name in self._whl_libraries:
-        diff = _diff_dict(self._whl_libraries[repo_name], repo.args)
+        self._whl_libraries[whl_repo_name] = {
+            k: v for k, v in repo.args.items()
+            if k not in forbidden_args | {
+                "config_load": True
+            }
+        }
+
+        args = repo.args
+
+        deps_args = {}
+        for key in ("config_load", "dep_template", "group_deps", "group_name", "annotation", "pip_data_exclude"):
+            if key in args and args[key] != None:
+                deps_args[key] = args[key]
+
+        deps_args["whl_library"] = "@{}//:BUILD.bazel".format(whl_repo_name)
+        repos_dict = self._whl_library_deps
+
+    repo_name = "{}_{}_{}".format(self.name, version_label(python_version), repo.repo_name)
+    if repo_name in repos_dict:
+        diff = _diff_dict(repos_dict[repo_name], deps_args)
         if diff:
             self._logger.fail(lambda: (
                 "Attempting to create a duplicate library {repo_name} for {whl_name} with different arguments. Already existing declaration has:\n".format(
@@ -331,7 +365,7 @@ def _add_whl_library(self, *, python_version, whl, repo):
                 ])
             ))
             return
-    self._whl_libraries[repo_name] = repo.args
+    repos_dict[repo_name] = deps_args
 
     mapping = self._whl_map.setdefault(whl.name, {})
     if repo.config_setting in mapping and mapping[repo.config_setting] != repo_name:
@@ -685,6 +719,7 @@ def _whl_repo(
 
     return struct(
         repo_name = whl_repo_name(src.filename, src.sha256, *target_platforms),
+        whl_repo_name = whl_repo_name(src.filename, src.sha256),
         args = args,
         config_setting = whl_config_setting(
             version = python_version,
diff --git a/python/private/pypi/labels.bzl b/python/private/pypi/labels.bzl
@@ -22,3 +22,5 @@ PY_LIBRARY_IMPL_LABEL = "_pkg"
 DATA_LABEL = "data"
 DIST_INFO_LABEL = "dist_info"
 NODEPS_LABEL = "no_deps"
+NODEPS_WHL_FILE_LABEL = "_whl_file"
+NODEPS_PY_LIBRARY_LABEL = "_srcs"
diff --git a/python/private/pypi/whl_extract.bzl b/python/private/pypi/whl_extract.bzl
@@ -26,6 +26,11 @@ def whl_extract(rctx, *, whl_path, logger):
         logger = logger,
     )
 
+    # Symlink the METADATA file to be able to refer to it from another repository_rule. This allows
+    # us to split the extracted sources and the closure itself to 2 different repository rules
+    # allowing us to not extract the same wheels multiple times.
+    rctx.symlink(metadata_file, "METADATA")
+
     # Get the <prefix>.dist_info dir name
     dist_info_dir = metadata_file.dirname
     rctx.file(
diff --git a/python/private/pypi/whl_library.bzl b/python/private/pypi/whl_library.bzl
diff --git a/python/private/pypi/whl_library_targets.bzl b/python/private/pypi/whl_library_targets.bzl
