perf(pypi): extract wheels once and reuse

python/private/pypi/extension.bzl

@@ -32,6 +32,10 @@ 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_archive.bzl", "whl_archive")
+load(":whl_deps_repo.bzl", "whl_deps_repo")
+load(":whl_repo_name.bzl", "whl_repo_name")
+load(":pep508_requirement.bzl", "requirement")
 
 def _whl_mods_impl(whl_mods_dict):
     """Implementation of the pip.whl_mods tag class.
@@ -537,6 +541,149 @@ def _create_unified_hub_repo(mods):
         packages = packages,
     )
 
+def register_whl_libraries(
+        registrations, 
+        rules = struct(
+            whl_library=whl_library,
+            whl_deps_repo=whl_deps_repo,
+            whl_archive=whl_archive,
+        ),
+    ):
+    """Register all of the whl libraries that where produced by the extension.
+
+    Args:
+        registrations: {type}`dict[str, dict[str, Any]]` The args to pass to
+            the repository rules when processing everything.
+        rules: {type}`struct` used to inject different implementations used to
+            create instances. Used for testing.
+    """
+    # we first check which wheels are ambiguous - belong to the same hub repository, indicated
+    # by the `dep_template`, but have different (URL, index_url) pairs for each whl.
+    urls = {}
+    fallback = {}
+    separate_srcs = {}
+    for name, args in registrations.items():
+        if "urls" not in args or "filename" not in args:
+            # only enabled with bazel downloader
+            continue
+
+        if "annotation" in args or "whl_patches" in args:
+            # No support for patched repos yet
+            continue
+
+        # For each each integrity, filename, we should have a unique url and index_url
+        # used across all of the hub repos, if not, print a warning and fallback to the
+        # previous behaviour of not using the source repository.
+        #
+        # If the filename is the same but it has a different sha256, then it means that
+        # there could be one of the following situations:
+        # 1. dependency confusion.
+        # 2. somebody patched the wheel.
+        # 3. it is an unlucky coincidence.
+        #
+        # Any of these cases are a bit tricky to get right, the repo name is derived
+        # filename together. The repo name is derived from the filename and a hash for
+        # this reason.
+        integrity = args["integrity"]
+        filename = args["filename"]
+
+        url = args["urls"][0]
+        index_url = args["index_url"]
+
+        hub_name, _, _ = args["dep_template"].partition("//")
+        hub_name = hub_name.strip("@")
+
+        # for now we include the hub_name in the key since the reuse is done for each
+        # hub_name separately.
+        key = (hub_name, filename, integrity)
+        value = (url, index_url)
+
+        if (filename, integrity) in fallback:
+            continue
+
+        got = urls.setdefault(key, value)
+        if got != value:
+            urls.pop(key)
+            fallback.setdefault((filename, integrity), None)
+
+        separate_srcs[name] = key
+
+    for name, args in registrations.items():
+        if name in separate_srcs:
+            # first handle the regular case
+            continue
+
+        rules.whl_library(name = name, **args)
+
+    # Then handle the reusable case
+    registered_srcs = {}
+    for name, key in separate_srcs.items():
+        args = registrations[name]
+        hub_name, _, integrity = key
+
+        filename = args["filename"]
+        extract_args = {
+            k: v
+            for k, v in args.items()
+            if k not in {
+                "config_load": None,
+                "dep_template": None,
+            }
+        }
+
+        # 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 = "{}_{}".format(
+            hub_name,
+            whl_repo_name(filename, ""),
+        )
+        if key not in registered_srcs:
+            registered_srcs[key] = extract_repo_name
+            rules.whl_archive(
+                name = extract_repo_name,
+                **extract_args,
+            )
+
+        deps_args = {
+            k: args.get(k)
+            for k in [
+                "config_load",
+                "requirement",
+                "dep_template",
+                "group_deps",
+                "group_name",
+                "repo_prefix",
+            ]
+            if args.get(k) != None
+        } | {
+            "metadata_file": "@{}//:metadata.json".format(extract_repo_name),
+        }
+        rules.whl_deps_repo(
+            name = name,
+            **deps_args
+        )
+
+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
+
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories.
 
@@ -611,8 +758,7 @@ def _pip_impl(module_ctx):
     # Build all of the wheel modifications if the tag class is called.
     _whl_mods_impl(mods.whl_mods)
 
-    for name, args in mods.whl_libraries.items():
-        whl_library(name = name, **args)
+    register_whl_libraries(registrations = mods.whl_libraries)
 
     for hub_name, whl_map in mods.hub_whl_map.items():
         hub_repository(

python/private/pypi/whl_library_deps_targets.bzl

@@ -23,8 +23,8 @@ def whl_library_deps_targets(
         repo,
         aliases = None,
         metadata_name,
-        requires_dist,
-        extras,
+        requires_dist = [],
+        extras = [],
         include = [],
         group_deps = [],
         group_name = None,
@@ -168,6 +168,7 @@ def whl_library_deps_targets(
                 package_deps = package_deps,
                 tmpl = dep_template.format(name = "{}", target = PY_LIBRARY_PUBLIC_LABEL),
             ),
+            precompile = "disabled",
             tags = tags,
             visibility = impl_vis,
         )

python/private/pypi/whl_library_targets.bzl

@@ -86,16 +86,17 @@ def whl_library_targets(
         **kwargs
     )
 
-    whl_library_deps_targets(
-        name = name,
-        metadata_name = metadata_name,
-        requires_dist = requires_dist,
-        dep_template = dep_template,  # only needed if requires_dist or group_name is present
-        group_deps = group_deps,  # only needed if group_name is present
-        group_name = group_name,  # must specify group_deps together
-        extras = extras,  # only needed if requires_dist is present
-        include = include,  # only needed if requires_dist is present
-        repo = None,  # set aliases in the same repo
-        aliases = {},
-        **kwargs
-    )
+    if dep_template:
+        whl_library_deps_targets(
+            name = name,
+            metadata_name = metadata_name,
+            requires_dist = requires_dist,
+            dep_template = dep_template,  # only needed if requires_dist or group_name is present
+            group_deps = group_deps,  # only needed if group_name is present
+            group_name = group_name,  # must specify group_deps together
+            extras = extras,  # only needed if requires_dist is present
+            include = include,  # only needed if requires_dist is present
+            repo = None,  # set aliases in the same repo
+            aliases = {},
+            **kwargs
+        )