feat(pypi): download and extract wheels once

Author: aignas

news/3978.fixed.md

@@ -0,0 +1,5 @@
+(pip) Wrap `pip_parse` and `unified_workspace_hub_repo` calls in the
+WORKSPACE file with `maybe` from `@bazel_tools//tools/build_defs/repo:utils.bzl`
+so that extracted wheel dependencies are reused and not re-created when the
+WORKSPACE is evaluated multiple times. This mirrors the `maybe` pattern already
+used in the bzlmod `pip` extension and in the generated `install_deps` macro.

python/private/BUILD.bazel

@@ -929,6 +929,12 @@ bzl_library(
     deps = [":visibility"],
 )
 
+bzl_library(
+    name = "repo_utils",
+    srcs = ["repo_utils.bzl"],
+    deps = [":text_util"],
+)
+
 bzl_library(
     name = "bzlmod_enabled",
     srcs = ["bzlmod_enabled.bzl"],
@@ -994,11 +1000,6 @@ bzl_library(
     srcs = ["py_runtime_info.bzl"],
 )
 
-bzl_library(
-    name = "repo_utils",
-    srcs = ["repo_utils.bzl"],
-)
-
 bzl_library(
     name = "runtimes_manifest_workspace",
     srcs = ["runtimes_manifest_workspace.bzl"],

python/private/internal_dev_deps.bzl

@@ -14,12 +14,15 @@
 """Module extension for internal dev_dependency=True setup."""
 
 load("@bazel_ci_rules//:rbe_repo.bzl", "rbe_preconfig")
-load("//python/private/pypi:whl_library.bzl", "whl_library")
+load("//python/private:repo_utils.bzl", "repo_utils")
+load("//python/private/pypi:whl_library.bzl", _whl_library = "whl_library")
 load("//tests/support/whl_from_dir:whl_from_dir_repo.bzl", "whl_from_dir_repo")
 load(":runtime_env_repo.bzl", "runtime_env_repo")
 
 def _internal_dev_deps_impl(mctx):
     _ = mctx  # @unused
+    maybe = repo_utils.maybe({})
+    whl_library = lambda **kwargs: _whl_library(maybe = maybe, **kwargs)
 
     # Creates a default toolchain config for RBE.
     # Use this as is if you are using the rbe_ubuntu16_04 container,
@@ -118,6 +121,7 @@ def _internal_dev_deps_impl(mctx):
         root = "//tests/pypi/whl_library/testdata/pkg:BUILD.bazel",
         output = "pkg-1.0-any-none-any.whl",
         requirement = "pkg[optional]",
+        maybe = maybe,
         # The following is necessary to enable pipstar and make tests faster
         config_load = "@rules_python//tests/pypi/whl_library/testdata:packages.bzl",
         dep_template = "@whl_library_extras_{name}//:{target}",
@@ -126,6 +130,7 @@ def _internal_dev_deps_impl(mctx):
         name = "whl_library_extras_optional_dep",
         root = "//tests/pypi/whl_library/testdata/optional_dep:BUILD.bazel",
         output = "optional_dep-1.0-any-none-any.whl",
+        maybe = maybe,
         requirement = "optional_dep",
         # The following is necessary to enable pipstar and make tests faster
         config_load = "@rules_python//tests/pypi/whl_library/testdata:packages.bzl",
@@ -149,15 +154,16 @@ def _internal_dev_deps_impl(mctx):
         },
     )
 
-def _whl_library_from_dir(*, name, output, root, **kwargs):
+def _whl_library_from_dir(*, name, output, root, maybe, **kwargs):
     whl_from_dir_repo(
         name = "{}_whl".format(name),
         root = root,
         output = output,
     )
-    whl_library(
+    _whl_library(
         name = name,
         whl_file = "@{}_whl//:{}".format(name, output),
+        maybe = maybe,
         **kwargs
     )
 

python/private/pypi/extension.bzl

@@ -22,6 +22,7 @@ load("//python/private:auth.bzl", "AUTH_ATTRS")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:pyproject_utils.bzl", "read_pyproject", "version_from_requires_python")
 load("//python/private:repo_utils.bzl", "repo_utils")
+load("//python/private:text_util.bzl", "render")
 load(":hub_builder.bzl", "hub_builder")
 load(":hub_repository.bzl", "hub_repository", "whl_config_settings_to_json")
 load(":parse_whl_name.bzl", "parse_whl_name")
@@ -463,10 +464,27 @@ You cannot use both the additive_build_content and additive_build_content_file a
         out = hub.build()
 
         for whl_name, lib in out.whl_libraries.items():
+            # NOTE @aignas 2026-07-04: if the same wheel is downloaded from multiple
+            # indexes, this will fail, forcing the user to actually download the wheel
+            # from the same and deterministic location. This is usually the case for
+            # public wheels and users should setup the defaults.index_url to correct
+            # fall-back in rules_python we should handle the default index to substitute
+            # any index-url in requirements pointing to the public PyPI mirrors.
             if whl_name in whl_libraries:
-                fail("'{}' already in created".format(whl_name))
-            else:
-                whl_libraries[whl_name] = lib
+                existing = whl_libraries[whl_name]
+
+                diff = repo_utils.diff_dict(existing, lib)
+                if diff:
+                    fail("'{}' already in created:\n{}".format(
+                        whl_name,
+                        "\n".join([
+                            "    {}: {}".format(key, render.indent(render.dict(value)).lstrip())
+                            for key, value in diff.items()
+                            if value
+                        ]),
+                    ))
+
+            whl_libraries[whl_name] = lib
 
         exposed_packages[hub.name] = out.exposed_packages
         extra_aliases[hub.name] = out.extra_aliases
@@ -611,8 +629,13 @@ def _pip_impl(module_ctx):
     # Build all of the wheel modifications if the tag class is called.
     _whl_mods_impl(mods.whl_mods)
 
+    registered = {}
     for name, args in mods.whl_libraries.items():
-        whl_library(name = name, **args)
+        whl_library(
+            name = name,
+            maybe = repo_utils.maybe(registered),
+            **args
+        )
 
     for hub_name, whl_map in mods.hub_whl_map.items():
         hub_repository(
@@ -991,7 +1014,7 @@ a string `"{os}_{arch}"` as the value here. You could also use `"{os}_{arch}_fre
 """,
         ),
         "uv_lock": attr.label(
-            doc = """
+            doc = """\
 (label, optional): A label pointing to the uv.lock file. If provided,
 the uv.lock file will be used as the primary source for package metadata.
 

python/private/pypi/hub_builder.bzl

@@ -87,6 +87,7 @@ def hub_builder(
         # Functions to download according to the config
         # dict[str python_version, callable]
         _get_index_urls = {},
+        _default_index_url = {},
         # Tells whether to use the downloader for a package.
         # dict[str python_version, dict[str package_name, bool use_downloader]]
         _use_downloader = {},
@@ -256,44 +257,6 @@ def _add_extra_aliases(self, extra_hub_aliases):
             {alias: True for alias in aliases},
         )
 
-def _diff_dict(first, second):
-    """A simple utility to shallow compare dictionaries.
-
-    Args:
-        first: The first dictionary to compare.
-        second: The second dictionary to compare.
-
-    Returns:
-        A dictionary containing the differences, with keys "common", "different",
-        "extra", and "missing", or None if the dictionaries are identical.
-    """
-    missing = {}
-    extra = {
-        key: value
-        for key, value in second.items()
-        if key not in first
-    }
-    common = {}
-    different = {}
-
-    for key, value in first.items():
-        if key not in second:
-            missing[key] = value
-        elif value == second[key]:
-            common[key] = value
-        else:
-            different[key] = (value, second[key])
-
-    if missing or extra or different:
-        return {
-            "common": common,
-            "different": different,
-            "extra": extra,
-            "missing": missing,
-        }
-    else:
-        return None
-
 def _add_whl_library(self, *, python_version, whl, repo):
     """Add a whl_library and kwargs to call it with for the hub.
 
@@ -309,16 +272,10 @@ def _add_whl_library(self, *, python_version, whl, repo):
         # disallow building from sdist.
         return
 
-    # 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 repo_name in self._whl_libraries:
-        diff = _diff_dict(self._whl_libraries[repo_name], repo.args)
+        diff = repo_utils.diff_dict(self._whl_libraries[repo_name], repo.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(
@@ -348,32 +305,33 @@ def _add_whl_library(self, *, python_version, whl, repo):
 ### end of setters, below we have various functions to implement the public methods
 
 def _set_get_index_urls(self, mctx, pip_attr):
+    python_version = pip_attr.python_version
+
     # Resolve the index URL through envsubst so the ``$VAR`` / ``${VAR:-default}``
     # form is honored when deciding whether the experimental index-url mode is
     # active. Without this, an unsubstituted template like ``$RULES_PYTHON_PIP_INDEX_URL``
     # is treated as truthy and the mode is forced on, even when the env var
     # would expand to the empty string.
-    default_index_url = envsubst(
+    self._default_index_url[python_version] = envsubst(
         pip_attr.experimental_index_url,
         pip_attr.envsubst,
         mctx.getenv,
     ) or self._config.index_url
     default_extra_index_urls = pip_attr.experimental_extra_index_urls or []
 
-    if not default_index_url:
+    if not self._default_index_url[python_version]:
         # parallel_download is set to True by default, so we are not checking/validating it
         # here
         return False
 
-    python_version = pip_attr.python_version
     self._use_downloader.setdefault(python_version, {}).update({
         normalize_name(s): False
         for s in pip_attr.simpleapi_skip
     })
     self._get_index_urls[python_version] = lambda ctx, distributions, *, index_url = None, extra_index_urls = None: self._simpleapi_download_fn(
         ctx,
         attr = struct(
-            index_url = (index_url or default_index_url).rstrip("/"),
+            index_url = (index_url or self._default_index_url[python_version]).rstrip("/"),
             extra_index_urls = [
                 x.rstrip("/")
                 for x in (extra_index_urls or default_extra_index_urls)

python/private/pypi/whl_library.bzl

@@ -14,6 +14,7 @@
 
 ""
 
+load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 load("//python/private:auth.bzl", "AUTH_ATTRS", "get_auth")
 load("//python/private:envsubst.bzl", "envsubst")
 load("//python/private:is_standalone_interpreter.bzl", "is_standalone_interpreter")
@@ -28,6 +29,7 @@ load(":pypi_repo_utils.bzl", "pypi_repo_utils")
 load(":urllib.bzl", "urllib")
 load(":whl_extract.bzl", "whl_extract")
 load(":whl_metadata.bzl", "parse_entry_points", "whl_metadata")
+load(":whl_repo_name.bzl", "whl_repo_name")
 
 _CPPFLAGS = "CPPFLAGS"
 _COMMAND_LINE_TOOLS_PATH_SLUG = "commandlinetools"
@@ -736,7 +738,7 @@ def _whl_deps_library_impl(rctx):
         repo = rctx.attr.repo or (
             str(rctx.attr.metadata_file) if rctx.attr.metadata_file else None
         ),
-        extras = requirement(rctx.attr.requirement).extras,
+        extras = rctx.attr.extras,
     )
     rctx.file("BUILD.bazel", build_file_contents)
 
@@ -748,9 +750,10 @@ whl_deps_library = repository_rule(
             "dep_template",
             "group_deps",
             "group_name",
-            "requirement",
+            "repo_prefix",
         ]
     } | {
+        "extras": attr.string_list(doc = "The extras to configure for this repo."),
         "metadata": attr.string(
             doc = """
 The subset of the METADATA contents that is needed for generation of the dependencies.
@@ -773,7 +776,7 @@ Does not depend on any python.
     environ = [REPO_DEBUG_ENV_VAR],
 )
 
-def whl_library(name, repo = None, **kwargs):
+def whl_library(name, repo = None, maybe = maybe, **kwargs):
     """Create a whl_library.
 
     This proxies to one of the underlying implementations:
@@ -783,6 +786,9 @@ def whl_library(name, repo = None, **kwargs):
     Args:
         name: {type}`str` The name of the repo.
         repo: Unused, will be dropped in the next major release.
+        maybe: This is the repo rule that is used in WORKSPACE mode and in the extension eval to
+            dedupe some of the invocations. This has to be overridden on bzlmod using
+            {obj}`repo_utils.bzlmod_maybe`.
         **kwargs: The args passed to the underlying implementation.
 
     Returns:
@@ -794,6 +800,80 @@ def whl_library(name, repo = None, **kwargs):
     urls = kwargs.get("urls", [])
     filename = kwargs.get("filename")
     if whl_file or (urls and filename and filename.endswith(".whl")):
-        whl_archive(name = name, **kwargs)
+        filename = filename or Label(whl_file).name
+        if "annotation" in kwargs or "whl_patches" in kwargs:
+            # No reuse of the whl_library because there is an annotation here
+            whl_archive(name = name, **kwargs)
+            return
+
+        extract_args = {
+            k: v
+            for k, v in kwargs.items()
+            if k not in {
+                # TODO @aignas 2026-08-01: what about python_interpreter and python_interpreter_target
+                "config_load": None,
+                "dep_template": None,
+            }
+        }
+        if "index_url" in extract_args:
+            # TODO @aignas 2026-08-01: figure out where we should do the fix here.
+            extract_args["index_url"] = extract_args["index_url"].strip("/")
+
+        # The extras do not affect the extraction, so normalize the requirement
+        # to allow the same wheel to be extracted only once.
+        extract_args["requirement"] = _without_extras(extract_args["requirement"])
+        extract_repo_name = "w_{}".format(
+            whl_repo_name(filename, kwargs.get("sha256")),
+        )
+
+        maybe(
+            whl_archive,
+            name = extract_repo_name,
+            **extract_args
+        )
+
+        req = requirement(kwargs["requirement"])
+
+        deps_args = {
+            k: kwargs.get(k)
+            for k in [
+                "config_load",
+                "dep_template",
+                "group_deps",
+                "group_name",
+            ]
+            if kwargs.get(k) != None
+        } | {
+            # TODO @aignas 2026-08-01: add extras only if the list is non-empty
+            "extras": req.extras,
+            "metadata_file": "@{}//:metadata.json".format(extract_repo_name),
+        }
+        whl_deps_library(
+            name = name,
+            **deps_args
+        )
     else:
+        # 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.
         pip_archive(name = name, **kwargs)
+
+def _without_extras(requirement_line):
+    """Remove the extras from a requirement line.
+
+    The extras do not affect which wheel is downloaded and extracted, so they
+    can be removed to allow the same wheel to be extracted only once even when
+    it is referenced with different extras.
+
+    Args:
+        requirement_line: {type}`str` the requirement line, e.g.
+            `foo[bar]==1.0 --hash=sha256:...`.
+
+    Returns:
+        The requirement line without the extras, e.g. `foo==1.0 --hash=sha256:...`.
+    """
+    name_and_version, _, extras = requirement_line.partition("[")
+    if not extras:
+        return requirement_line
+    _, _, rest = extras.partition("]")
+    return name_and_version + rest