perf(pypi): extract wheels once and reuse

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/BUILD.bazel

@@ -473,6 +473,7 @@ bzl_library(
         ":urllib",
         ":whl_extract",
         ":whl_metadata",
+        ":whl_repo_name",
         "//python/private:auth",
         "//python/private:envsubst",
         "//python/private:is_standalone_interpreter",
@@ -522,6 +523,12 @@ bzl_library(
     deps = ["//python/private:envsubst"],
 )
 
+bzl_library(
+    name = "index_sources",
+    srcs = ["index_sources.bzl"],
+    deps = [":hash"],
+)
+
 bzl_library(
     name = "argparse",
     srcs = ["argparse.bzl"],
@@ -542,12 +549,6 @@ bzl_library(
     srcs = ["hash.bzl"],
 )
 
-bzl_library(
-    name = "index_sources",
-    srcs = ["index_sources.bzl"],
-    deps = [":hash"],
-)
-
 bzl_library(
     name = "labels",
     srcs = ["labels.bzl"],

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

@@ -88,6 +88,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 = {},
@@ -257,44 +258,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.
 
@@ -310,16 +273,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(
@@ -349,32 +306,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)