feat(bzlmod): support patching 'whl' distributions

.bazelrc

@@ -3,8 +3,8 @@
 # This lets us glob() up all the files inside the examples to make them inputs to tests
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, run tools/bazel_integration_test/update_deleted_packages.sh
-build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
 
 test --test_output=errors
 

CHANGELOG.md

@@ -46,6 +46,11 @@ Breaking changes:
 * (py_wheel) Produce deterministic wheel files and make `RECORD` file entries
   follow the order of files written to the `.whl` archive.
 
+### Added
+
+* (bzlmod) Added `.whl` patching support via `patches` and `patch_strip`
+  arguments to the new `pip.override` tag class.
+
 ## [0.26.0] - 2023-10-06
 
 ### Changed

examples/bzlmod/MODULE.bazel

@@ -113,6 +113,19 @@ pip.parse(
         "@whl_mods_hub//:wheel.json": "wheel",
     },
 )
+
+# You can add patches that will be applied on the whl contents.
+#
+# The patches have to be in the unified-diff format.
+pip.override(
+    file = "requests-2.25.1-py2.py3-none-any.whl",
+    patch_strip = 1,
+    patches = [
+        "@//patches:empty.patch",
+        "@//patches:requests_metadata.patch",
+        "@//patches:requests_record.patch",
+    ],
+)
 use_repo(pip, "pip")
 
 bazel_dep(name = "other_module", version = "", repo_name = "our_other_module")

examples/bzlmod/patches/BUILD.bazel

@@ -0,0 +1,4 @@
+exports_files(
+    srcs = glob(["*.patch"]),
+    visibility = ["//visibility:public"],
+)

examples/bzlmod/patches/requests_metadata.patch

@@ -0,0 +1,12 @@
+diff --unified --recursive a/requests-2.25.1.dist-info/METADATA b/requests-2.25.1.dist-info/METADATA
+--- a/requests-2.25.1.dist-info/METADATA	2020-12-16 19:37:50.000000000 +0900
++++ b/requests-2.25.1.dist-info/METADATA	2023-09-30 20:31:50.079863410 +0900
+@@ -1,7 +1,7 @@
+ Metadata-Version: 2.1
+ Name: requests
+ Version: 2.25.1
+-Summary: Python HTTP for Humans.
++Summary: Python HTTP for Humans. Patched.
+ Home-page: https://requests.readthedocs.io
+ Author: Kenneth Reitz
+ Author-email: [email protected]

examples/bzlmod/patches/requests_record.patch

@@ -0,0 +1,11 @@
+--- a/requests-2.25.1.dist-info/RECORD
++++ b/requests-2.25.1.dist-info/RECORD
+@@ -17,7 +17,7 @@
+ requests/structures.py,sha256=msAtr9mq1JxHd-JRyiILfdFlpbJwvvFuP3rfUQT_QxE,3005
+ requests/utils.py,sha256=_K9AgkN6efPe-a-zgZurXzds5PBC0CzDkyjAE2oCQFQ,30529
+ requests-2.25.1.dist-info/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
+-requests-2.25.1.dist-info/METADATA,sha256=RuNh38uN0IMsRT3OwaTNB_WyGx6RMwwQoMwujXfkUVM,4168
++requests-2.25.1.dist-info/METADATA,sha256=fRSAA0u0Bi0heD4zYq91wdNUTJlbzhK6_iDOcRRNDx4,4177
+ requests-2.25.1.dist-info/WHEEL,sha256=Z-nyYpwrcSqxfdux5Mbn_DQ525iP7J2DG3JgGvOYyTQ,110
+ requests-2.25.1.dist-info/top_level.txt,sha256=fMSVmHfb5rbGOo6xv-O_tUX6j-WyixssE-SnwcDRxNQ,9
+ requests-2.25.1.dist-info/RECORD,,

examples/bzlmod/whl_mods/appended_build_content.BUILD

@@ -5,3 +5,12 @@ write_file(
     out = "generated_file.txt",
     content = ["Hello world from requests"],
 )
+
+filegroup(
+    name = "whl_orig",
+    srcs = glob(
+        ["*.whl"],
+        allow_empty = False,
+        exclude = ["*-patched-*.whl"],
+    ),
+)

python/pip_install/BUILD.bazel

@@ -30,6 +30,7 @@ bzl_library(
         "//python/pip_install/private:srcs_bzl",
         "//python/private:bzlmod_enabled_bzl",
         "//python/private:normalize_name_bzl",
+        "//python/private:patch_whl_bzl",
         "//python/private:render_pkg_aliases_bzl",
         "//python/private:toolchains_repo_bzl",
         "//python/private:which_bzl",
@@ -97,6 +98,8 @@ filegroup(
     srcs = [
         "//python/pip_install/tools/dependency_resolver:py_srcs",
         "//python/pip_install/tools/wheel_installer:py_srcs",
+        "//python/private:repack_whl.py",
+        "//tools:wheelmaker.py",
     ],
     visibility = ["//python/pip_install/private:__pkg__"],
 )

python/pip_install/pip_repository.bzl

@@ -22,6 +22,7 @@ load("//python/pip_install/private:generate_whl_library_build_bazel.bzl", "gener
 load("//python/pip_install/private:srcs.bzl", "PIP_INSTALL_PY_SRCS")
 load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")
 load("//python/private:normalize_name.bzl", "normalize_name")
+load("//python/private:patch_whl.bzl", "patch_whl")
 load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
 load("//python/private:toolchains_repo.bzl", "get_host_os_arch")
 load("//python/private:which.bzl", "which_with_fail")
@@ -44,6 +45,7 @@ def _construct_pypath(rctx):
 
     Args:
         rctx: Handle to the repository_context.
+
     Returns: String of the PYTHONPATH.
     """
 
@@ -542,6 +544,22 @@ def _whl_library_impl(rctx):
     if not rctx.delete("whl_file.json"):
         fail("failed to delete the whl_file.json file")
 
+    if rctx.attr.whl_patches:
+        patches = {}
+        for patch_file, json_args in patches.items():
+            patch_dst = struct(**json.decode(json_args))
+            if whl_path.basename in patch_dst.whls:
+                patches[patch_file] = patch_dst.patch_strip
+
+        whl_path = patch_whl(
+            rctx,
+            python_interpreter = python_interpreter,
+            whl_path = whl_path,
+            patches = patches,
+            quiet = rctx.attr.quiet,
+            timeout = rctx.attr.timeout,
+        )
+
     result = rctx.execute(
         args + ["--whl-file", whl_path],
         environment = environment,
@@ -635,6 +653,13 @@ whl_library_attrs = {
         mandatory = True,
         doc = "Python requirement string describing the package to make available",
     ),
+    "whl_patches": attr.label_keyed_string_dict(
+        doc = """"a label-keyed-string dict that has
+            json.encode(struct([whl_file], patch_strip]) as values. This
+            is to maintain flexibility and correct bzlmod extension interface
+            until we have a better way to define whl_library and move whl
+            patching to a separate place. INTERNAL USE ONLY.""",
+    ),
     "_python_path_entries": attr.label_list(
         # Get the root directory of these rules and keep them as a default attribute
         # in order to avoid unnecessary repository fetching restarts.

python/pip_install/private/srcs.bzl

@@ -13,4 +13,6 @@ PIP_INSTALL_PY_SRCS = [
     "@rules_python//python/pip_install/tools/wheel_installer:namespace_pkgs.py",
     "@rules_python//python/pip_install/tools/wheel_installer:wheel.py",
     "@rules_python//python/pip_install/tools/wheel_installer:wheel_installer.py",
+    "@rules_python//python/private:repack_whl.py",
+    "@rules_python//tools:wheelmaker.py",
 ]

python/private/BUILD.bazel

@@ -88,6 +88,17 @@ bzl_library(
     srcs = ["normalize_name.bzl"],
 )
 
+bzl_library(
+    name = "patch_whl_bzl",
+    srcs = ["patch_whl.bzl"],
+    deps = [":parse_whl_name_bzl"],
+)
+
+bzl_library(
+    name = "parse_whl_name_bzl",
+    srcs = ["parse_whl_name.bzl"],
+)
+
 bzl_library(
     name = "py_cc_toolchain_bzl",
     srcs = [
@@ -239,13 +250,14 @@ bzl_library(
 exports_files(
     [
         "coverage.patch",
+        "repack_whl.py",
+        "py_cc_toolchain_rule.bzl",
         "py_package.bzl",
         "py_wheel.bzl",
         "py_wheel_normalize_pep440.bzl",
         "reexports.bzl",
         "stamp.bzl",
         "util.bzl",
-        "py_cc_toolchain_rule.bzl",
     ],
     visibility = ["//:__subpackages__"],
 )

python/private/bzlmod/pip.bzl

@@ -26,6 +26,7 @@ load(
 load("//python/pip_install:requirements_parser.bzl", parse_requirements = "parse")
 load("//python/private:full_version.bzl", "full_version")
 load("//python/private:normalize_name.bzl", "normalize_name")
+load("//python/private:parse_whl_name.bzl", "parse_whl_name")
 load("//python/private:version_label.bzl", "version_label")
 load(":pip_repository.bzl", "pip_repository")
 
@@ -78,7 +79,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
             whl_mods = whl_mods,
         )
 
-def _create_whl_repos(module_ctx, pip_attr, whl_map):
+def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
     python_interpreter_target = pip_attr.python_interpreter_target
 
     # if we do not have the python_interpreter set in the attributes
@@ -131,6 +132,10 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map):
             repo = pip_name,
             repo_prefix = pip_name + "_",
             annotation = annotation,
+            whl_patches = {
+                p: json.encode(args)
+                for p, args in whl_overrides.get(whl_name, {}).items()
+            },
             python_interpreter = pip_attr.python_interpreter,
             python_interpreter_target = python_interpreter_target,
             quiet = pip_attr.quiet,
@@ -217,6 +222,35 @@ def _pip_impl(module_ctx):
     # Build all of the wheel modifications if the tag class is called.
     _whl_mods_impl(module_ctx)
 
+    _overriden_whl_set = {}
+    whl_overrides = {}
+
+    for module in module_ctx.modules:
+        for attr in module.tags.override:
+            if not module.is_root:
+                fail("overrides are only supported in root modules")
+
+            if not attr.file.endswith(".whl"):
+                fail("Only whl overrides are supported at this time")
+
+            whl_name = normalize_name(parse_whl_name(attr.file).distribution)
+
+            if attr.file in _overriden_whl_set:
+                fail("Duplicate module overrides for '{}'".format(attr.file))
+            _overriden_whl_set[attr.file] = None
+
+            for patch in attr.patches:
+                if whl_name not in whl_overrides:
+                    whl_overrides[whl_name] = {}
+
+                if patch not in whl_overrides[whl_name]:
+                    whl_overrides[whl_name][patch] = struct(
+                        patch_strip = attr.patch_strip,
+                        whls = [],
+                    )
+
+                whl_overrides[whl_name][patch].whls.append(attr.file)
+
     # Used to track all the different pip hubs and the spoke pip Python
     # versions.
     pip_hub_map = {}
@@ -261,7 +295,7 @@ def _pip_impl(module_ctx):
             else:
                 pip_hub_map[pip_attr.hub_name].python_versions.append(pip_attr.python_version)
 
-            _create_whl_repos(module_ctx, pip_attr, hub_whl_map)
+            _create_whl_repos(module_ctx, pip_attr, hub_whl_map, whl_overrides)
 
     for hub_name, whl_map in hub_whl_map.items():
         pip_repository(
@@ -381,6 +415,35 @@ cannot have a child module that uses the same `hub_name`.
     }
     return attrs
 
+# NOTE: the naming of 'override' is taken from the bzlmod native
+# 'archive_override', 'git_override' bzlmod functions.
+_override_tag = tag_class(
+    attrs = {
+        "file": attr.string(
+            doc = """\
+The Python distribution file name which needs to be patched. This will be
+applied to all repositories that setup this distribution via the pip.parse tag
+class.""",
+            mandatory = True,
+        ),
+        "patch_strip": attr.int(
+            default = 0,
+            doc = """\
+The number of leading path segments to be stripped from the file name in the
+patches.""",
+        ),
+        "patches": attr.label_list(
+            doc = """\
+A list of patches to apply to the repository *after* 'whl_library' is extracted
+and BUILD.bazel file is generated.""",
+            mandatory = True,
+        ),
+    },
+    doc = """\
+Apply any overrides (e.g. patches) to a given Python distribution defined by
+other tags in this extension.""",
+)
+
 def _extension_extra_args():
     args = {}
 
@@ -412,6 +475,7 @@ the BUILD files for wheels.
 """,
     implementation = _pip_impl,
     tag_classes = {
+        "override": _override_tag,
         "parse": tag_class(
             attrs = _pip_parse_ext_attrs(),
             doc = """\

python/private/parse_whl_name.bzl

@@ -0,0 +1,72 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+A starlark implementation of a Wheel filename parsing.
+"""
+
+def parse_whl_name(file):
+    """Parse whl file name into a struct of constituents.
+
+    Args:
+        file (str): The file name of a wheel
+
+    Returns:
+        A struct with the following attributes:
+            distribution: the distribution name
+            version: the version of the distribution
+            build_tag: the build tag for the wheel. None if there was no
+              build_tag in the given string.
+            python_tag: the python tag for the wheel
+            abi_tag: the ABI tag for the wheel
+            platform_tag: the platform tag
+    """
+    if not file.endswith(".whl"):
+        fail("not a valid wheel: {}".format(file))
+
+    file = file[:-len(".whl")]
+
+    # Parse the following
+    # {distribution}-{version}(-{build tag})?-{python tag}-{abi tag}-{platform tag}.whl
+    #
+    # For more info, see the following standards:
+    # https://packaging.python.org/en/latest/specifications/binary-distribution-format/#binary-distribution-format
+    # https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/
+    head, _, platform_tag = file.rpartition("-")
+    if not platform_tag:
+        fail("cannot extract platform tag from the whl filename: {}".format(file))
+    head, _, abi_tag = head.rpartition("-")
+    if not abi_tag:
+        fail("cannot extract abi tag from the whl filename: {}".format(file))
+    head, _, python_tag = head.rpartition("-")
+    if not python_tag:
+        fail("cannot extract python tag from the whl filename: {}".format(file))
+    head, _, version = head.rpartition("-")
+    if not version:
+        fail("cannot extract version from the whl filename: {}".format(file))
+    distribution, _, maybe_version = head.partition("-")
+
+    if maybe_version:
+        version, build_tag = maybe_version, version
+    else:
+        build_tag = None
+
+    return struct(
+        distribution = distribution,
+        version = version,
+        build_tag = build_tag,
+        python_tag = python_tag,
+        abi_tag = abi_tag,
+        platform_tag = platform_tag,
+    )