fix(pypi): fetch retained wheel archives lazily

examples/bzlmod/whl_mods/BUILD.bazel

@@ -8,8 +8,10 @@ exports_files(
 py_test(
     name = "pip_whl_mods_test",
     srcs = ["pip_whl_mods_test.py"],
+    data = ["@pip//requests:whl"],
     env = {
         "REQUESTS_PKG": "$(rlocationpaths @pip//requests:pkg)",
+        "REQUESTS_WHL": "$(rlocationpath @pip//requests:whl)",
         "WHEEL_PKG": "$(rlocationpaths @pip//wheel:pkg)",
     },
     main = "pip_whl_mods_test.py",

examples/bzlmod/whl_mods/appended_build_content.BUILD

@@ -7,10 +7,8 @@ write_file(
 )
 
 filegroup(
-    name = "whl_orig",
-    srcs = glob(
-        ["*.whl"],
-        allow_empty = False,
-        exclude = ["*-patched-*.whl"],
-    ),
+    name = "whl_archive",
+    # The archive may live in a separate repository. Use the wheel target
+    # instead of assuming that the archive is beside the extracted files.
+    srcs = [":whl_file"],
 )

examples/bzlmod/whl_mods/pip_whl_mods_test.py

@@ -19,6 +19,7 @@
 import subprocess
 import sys
 import unittest
+import zipfile
 from pathlib import Path
 
 from python.runfiles import runfiles
@@ -145,6 +146,14 @@ def test_patches(self):
             "Summary: Python HTTP for Humans. Patched.",
             metadata.read_text().splitlines(),
         )
+        # The public wheel target must expose the same patches as the library,
+        # even when its archive is stored outside the extracted repository.
+        archive = Path(r.Rlocation(os.environ["REQUESTS_WHL"]))
+        with zipfile.ZipFile(archive) as wheel:
+            self.assertEqual(
+                wheel.read(f"requests-{current_wheel_version}.dist-info/METADATA"),
+                metadata.read_bytes(),
+            )
 
 
 if __name__ == "__main__":

news/lazy-wheel-archives.changed.md

@@ -0,0 +1,6 @@
+(pypi) Downloaded wheels are removed after extraction when their original
+archive can be fetched independently. The existing `whl` and `whl_file` targets
+fetch the archive from a separate repository on demand. Matching patches are
+applied in the separate wheel-file repository, which also supplies the patched
+archive for extraction. Wheels built from source retain their existing archive
+behavior.

python/private/pypi/BUILD.bazel

@@ -479,8 +479,10 @@ bzl_library(
     name = "whl_library",
     srcs = ["whl_library.bzl"],
     deps = [
+        ":patch_whl",
         ":pip_archive",
         ":whl_archive",
+        ":whl_file_repo",
     ],
 )
 
@@ -555,8 +557,20 @@ bzl_library(
     deps = [
         ":attrs",
         ":patch_and_extract_whl",
-        ":urllib",
         ":whl_deps_repo",
+        ":whl_file_repo",
+        "//python/private:auth",
+        "//python/private:repo_utils",
+    ],
+)
+
+bzl_library(
+    name = "whl_file_repo",
+    srcs = ["whl_file_repo.bzl"],
+    deps = [
+        ":attrs",
+        ":patch_whl",
+        ":urllib",
         "//python/private:auth",
         "//python/private:repo_utils",
     ],

python/private/pypi/patch_and_extract_whl.bzl

@@ -3,7 +3,7 @@
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:repo_utils.bzl", "repo_utils")
 load(":generate_whl_library_build_bazel.bzl", "generate_whl_library_build_bazel")
-load(":patch_whl.bzl", "patch_whl")
+load(":patch_whl.bzl", "get_whl_patches", "patch_whl")
 load(":pep508_requirement.bzl", "requirement")
 load(":pypi_repo_utils.bzl", "pypi_repo_utils")
 load(":whl_extract.bzl", "whl_extract")
@@ -69,7 +69,7 @@ def _remove_files(rctx, *basenames):
         elif path.is_dir:
             paths.extend(path.readdir())
 
-def patch_and_extract_whl(rctx, *, whl_path, logger, sdist_filename = None):
+def patch_and_extract_whl(rctx, *, whl_path, logger, sdist_filename = None, whl_file = None):
     """Extract the wheel, apply patches and generate BUILD.bazel files.
 
     Reused in pip and http wheel download code.
@@ -79,27 +79,30 @@ def patch_and_extract_whl(rctx, *, whl_path, logger, sdist_filename = None):
         whl_path: the whl path to extract.
         logger: The logger to use
         sdist_filename: The filename to ignore in the BUILD.bazel files as sources.
+        whl_file: The label of an externally stored wheel, if any.
 
     Returns:
         The repository metadata if the extraction is reproducible
     """
     if rctx.attr.whl_patches:
-        patches = {}
-        for patch_file, json_args in rctx.attr.whl_patches.items():
-            patch_dst = struct(**json.decode(json_args))
-            if whl_path.basename in patch_dst.whls:
-                patches[patch_file] = patch_dst.patch_strip
+        patches = get_whl_patches(rctx.attr.whl_patches, whl_path.basename)
 
         if patches:
+            # patch_whl deletes its input before repacking. Never delete an
+            # archive owned by another repository.
+            if whl_path.dirname != rctx.path("."):
+                rctx.symlink(whl_path, whl_path.basename)
+                whl_path = rctx.path(whl_path.basename)
             whl_path = patch_whl(
                 rctx,
                 whl_path = whl_path,
                 patches = patches,
             )
+            whl_file = None
 
     whl_extract(rctx, whl_path = whl_path, logger = logger)
 
-    install_dir_path = whl_path.dirname.get_child("site-packages")
+    install_dir_path = rctx.path("site-packages")
     metadata = whl_metadata(
         install_dir = install_dir_path,
         read_fn = rctx.read,
@@ -139,8 +142,14 @@ def patch_and_extract_whl(rctx, *, whl_path, logger, sdist_filename = None):
             metadata = metadata,
             filename = sdist_filename or whl_path.basename,
         ),
+        **({"whl_file": str(whl_file)} if whl_file else {})
     )
 
+    # The generated wheel target can obtain the original archive independently.
+    # Do not retain the temporary download or delete a caller-owned wheel.
+    if whl_file and whl_path.dirname == rctx.path("."):
+        rctx.delete(whl_path)
+
     # Delete these in case the wheel had them. They generally don't cause
     # a problem, but let's avoid the chance of that happening.
     rctx.file("WORKSPACE")

python/private/pypi/patch_whl.bzl

@@ -58,6 +58,23 @@ def patched_whl_name(original_whl_name):
         platform_tag = parsed_whl.platform_tag,
     )
 
+def get_whl_patches(whl_patches, filename):
+    """Select the patches that apply to a wheel filename.
+
+    Args:
+        whl_patches: Patch labels mapped to JSON-encoded wheel names and strip levels.
+        filename: The wheel filename to match.
+
+    Returns:
+        A dictionary of matching patch labels mapped to strip levels.
+    """
+    patches = {}
+    for patch_file, json_args in whl_patches.items():
+        patch_dst = struct(**json.decode(json_args))
+        if filename in patch_dst.whls:
+            patches[patch_file] = patch_dst.patch_strip
+    return patches
+
 def patch_whl(rctx, *, whl_path, patches):
     """Patch a whl file and repack it natively.
 

python/private/pypi/whl_archive.bzl

@@ -1,11 +1,11 @@
 ""
 
-load("//python/private:auth.bzl", "AUTH_ATTRS", "get_auth")
+load("//python/private:auth.bzl", "AUTH_ATTRS")
 load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
 load(":attrs.bzl", "ATTRS")
 load(":patch_and_extract_whl.bzl", "patch_and_extract_whl")
-load(":urllib.bzl", "urllib")
 load(":whl_deps_repo.bzl", "whl_deps_attrs")
+load(":whl_file_repo.bzl", "download_whl")
 
 def _whl_archive_impl(rctx):
     logger = repo_utils.logger(rctx)
@@ -14,47 +14,17 @@ def _whl_archive_impl(rctx):
     if rctx.attr.whl_file:
         rctx.watch(rctx.attr.whl_file)
         whl_path = rctx.path(rctx.attr.whl_file)
-
-        # Simulate the behaviour where the whl is present in the current directory.
-        rctx.symlink(whl_path, whl_path.basename)
-        whl_path = rctx.path(whl_path.basename)
     elif rctx.attr.urls and rctx.attr.filename:
-        filename = rctx.attr.filename
-        urls = rctx.attr.urls
-        urls = [
-            urllib.absolute_url(
-                rctx.attr.index_url,
-                url,
-                envsubst = rctx.attr.envsubst,
-                getenv = rctx.getenv,
-            )
-            for url in urls
-        ]
-        result = rctx.download(
-            url = urls,
-            output = filename,
-            sha256 = rctx.attr.sha256,
-            integrity = rctx.attr.integrity if not rctx.attr.sha256 else "",
-            auth = get_auth(rctx, urls),
-        )
-        if not rctx.attr.sha256 and not rctx.attr.integrity:
-            # this is only seen when there is a direct URL reference without a hash
-            logger.warn("Please update the requirement line to include the hash:\n{} \\\n    --hash=sha256:{}".format(
-                rctx.attr.requirement,
-                result.sha256,
-            ))
-
-        if not result.success:
-            fail("could not download the '{}' from {}:\n{}".format(filename, urls, result))
-
-        if filename.endswith(".whl"):
-            whl_path = rctx.path(filename)
-        else:
-            fail("Only wheels are supported")
+        whl_path = download_whl(rctx, logger = logger)
     else:
         fail("Either 'whl_file' or 'urls' and 'filename' needs to be specified")
 
-    return patch_and_extract_whl(rctx, whl_path = whl_path, logger = logger)
+    return patch_and_extract_whl(
+        rctx,
+        whl_path = whl_path,
+        whl_file = rctx.attr.whl_file or rctx.attr.downloaded_whl,
+        logger = logger,
+    )
 
 whl_archive_attrs = whl_deps_attrs | {
     "annotation": attr.label(
@@ -64,6 +34,9 @@ whl_archive_attrs = whl_deps_attrs | {
         ),
         allow_files = True,
     ),
+    "downloaded_whl": attr.string(
+        doc = "Label of the lazy archive target, emitted in BUILD.bazel without fetching it. INTERNAL USE ONLY.",
+    ),
     "filename": attr.string(
         doc = "Download the whl file to this filename. Only used when the `urls` is passed. If not specified, will be auto-detected from the `urls`.",
     ),

python/private/pypi/whl_file_repo.bzl

@@ -0,0 +1,93 @@
+"""Store an original or patched wheel separately from its extracted files."""
+
+load("//python/private:auth.bzl", "AUTH_ATTRS", "get_auth")
+load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
+load(":attrs.bzl", "ATTRS")
+load(":patch_whl.bzl", "patch_whl")
+load(":urllib.bzl", "urllib")
+
+def download_whl(rctx, *, logger):
+    """Download a wheel and return its path, preserving index authentication.
+
+    Args:
+        rctx: The repository context.
+        logger: The repository logger.
+
+    Returns:
+        The path of the downloaded wheel.
+    """
+    filename = rctx.attr.filename
+    if not filename.endswith(".whl"):
+        fail("Only wheels are supported")
+    urls = [
+        urllib.absolute_url(
+            rctx.attr.index_url,
+            url,
+            envsubst = rctx.attr.envsubst,
+            getenv = rctx.getenv,
+        )
+        for url in rctx.attr.urls
+    ]
+    result = rctx.download(
+        url = urls,
+        output = filename,
+        sha256 = rctx.attr.sha256,
+        integrity = rctx.attr.integrity if not rctx.attr.sha256 else "",
+        auth = get_auth(rctx, urls),
+    )
+    if not rctx.attr.sha256 and not rctx.attr.integrity:
+        logger.warn("Please update the requirement line to include the hash:\n{} \\\n    --hash=sha256:{}".format(
+            rctx.attr.requirement,
+            result.sha256,
+        ))
+    if not result.success:
+        fail("could not download the '{}' from {}:\n{}".format(filename, urls, result))
+    return rctx.path(filename)
+
+whl_file_repo_attrs = {
+    "envsubst": ATTRS["envsubst"],
+    "filename": attr.string(doc = "The filename of the downloaded wheel."),
+    "index_url": attr.string(doc = "The index URL used to resolve relative wheel URLs."),
+    "integrity": attr.string(doc = "The expected SRI checksum, used when sha256 is unset."),
+    "patches": attr.label_keyed_string_dict(doc = "Patch labels mapped to their strip levels."),
+    "requirement": attr.string(doc = "The requirement used in download diagnostics."),
+    "sha256": attr.string(doc = "The expected SHA256 checksum."),
+    "urls": attr.string_list(doc = "The wheel download URLs."),
+    "whl_file": attr.label(doc = "A caller-owned wheel to use instead of downloading."),
+} | AUTH_ATTRS
+
+def _whl_file_repo_impl(rctx):
+    if rctx.attr.whl_file:
+        rctx.watch(rctx.attr.whl_file)
+        whl_path = rctx.path(rctx.attr.whl_file)
+
+        # patch_whl deletes its input before repacking. Keep caller-owned
+        # wheels intact by giving it a repository-local symlink.
+        rctx.symlink(whl_path, whl_path.basename)
+        whl_path = rctx.path(whl_path.basename)
+    else:
+        whl_path = download_whl(rctx, logger = repo_utils.logger(rctx))
+    if rctx.attr.patches:
+        whl_path = patch_whl(
+            rctx,
+            whl_path = whl_path,
+            patches = {patch: int(strip) for patch, strip in rctx.attr.patches.items()},
+        )
+
+        # Patching temporarily unpacks the input. Retain only the replacement
+        # wheel; the installation is extracted in the consuming repository.
+        for path in rctx.path(".").readdir():
+            if path.basename != whl_path.basename:
+                rctx.delete(path)
+    rctx.file("BUILD.bazel", "exports_files([{}])\n".format(repr(whl_path.basename)))
+    if hasattr(rctx, "repo_metadata"):
+        return rctx.repo_metadata(reproducible = True)
+    return None
+
+whl_file_repo = repository_rule(
+    implementation = _whl_file_repo_impl,
+    attrs = whl_file_repo_attrs | {
+        "_rule_name": attr.string(default = "whl_file_repo"),
+    },
+    environ = [REPO_DEBUG_ENV_VAR],
+)