feat: reuse extracted whl srcs within a hub repository

This is the common denominator and for now no warnings are printed, but
there are opportunities to do this.

This approach is way more surgical than the previous one.

Fixes #2948

Author: aignas

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
+        )