pycross: Add patching support to py_wheel_library (#1436)

This patch adds a few arguments to `py_wheel_library` to simulate how
`http_archive` accepts patch-related arguments.

I also amended the existing test to validate the behaviour at a very
high level.

References: #1360

6 files changed, 182 insertions, 4 deletions

diff --git a/tests/pycross/0001-Add-new-file-for-testing-patch-support.patch b/tests/pycross/0001-Add-new-file-for-testing-patch-support.patch
new file mode 100644
index 0000000..fcbc309
--- /dev/null
+++ b/tests/pycross/0001-Add-new-file-for-testing-patch-support.patch
@@ -0,0 +1,17 @@
+From b2ebe6fe67ff48edaf2ae937d24b1f0b67c16f81 Mon Sep 17 00:00:00 2001
+From: Philipp Schrader <philipp.schrader@gmail.com>
+Date: Thu, 28 Sep 2023 09:02:44 -0700
+Subject: [PATCH] Add new file for testing patch support
+
+---
+ site-packages/numpy/file_added_via_patch.txt | 1 +
+ 1 file changed, 1 insertion(+)
+ create mode 100644 site-packages/numpy/file_added_via_patch.txt
+
+diff --git a/site-packages/numpy/file_added_via_patch.txt b/site-packages/numpy/file_added_via_patch.txt
+new file mode 100644
+index 0000000..9d947a4
+--- /dev/null
++++ b/site-packages/numpy/file_added_via_patch.txt
+@@ -0,0 +1 @@
++Hello from a patch!
diff --git a/tests/pycross/BUILD.bazel b/tests/pycross/BUILD.bazel
index 4f01272..52d1d18 100644
--- a/tests/pycross/BUILD.bazel
+++ b/tests/pycross/BUILD.bazel
@@ -32,3 +32,33 @@ py_test(
         "//python/runfiles",
     ],
 )
+
+py_wheel_library(
+    name = "patched_extracted_wheel_for_testing",
+    patch_args = [
+        "-p1",
+    ],
+    patch_tool = "patch",
+    patches = [
+        "0001-Add-new-file-for-testing-patch-support.patch",
+    ],
+    target_compatible_with = select({
+        # We don't have `patch` available on the Windows CI machines.
+        "@platforms//os:windows": ["@platforms//:incompatible"],
+        "//conditions:default": [],
+    }),
+    wheel = "@wheel_for_testing//file",
+)
+
+py_test(
+    name = "patched_py_wheel_library_test",
+    srcs = [
+        "patched_py_wheel_library_test.py",
+    ],
+    data = [
+        ":patched_extracted_wheel_for_testing",
+    ],
+    deps = [
+        "//python/runfiles",
+    ],
+)
diff --git a/tests/pycross/patched_py_wheel_library_test.py b/tests/pycross/patched_py_wheel_library_test.py
new file mode 100644
index 0000000..4591187
--- /dev/null
+++ b/tests/pycross/patched_py_wheel_library_test.py
@@ -0,0 +1,38 @@
+# 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.
+
+import unittest
+from pathlib import Path
+
+from python.runfiles import runfiles
+
+RUNFILES = runfiles.Create()
+
+
+class TestPyWheelLibrary(unittest.TestCase):
+    def setUp(self):
+        self.extraction_dir = Path(
+            RUNFILES.Rlocation("rules_python/tests/pycross/patched_extracted_wheel_for_testing")
+        )
+        self.assertTrue(self.extraction_dir.exists(), self.extraction_dir)
+        self.assertTrue(self.extraction_dir.is_dir(), self.extraction_dir)
+
+    def test_patched_file_contents(self):
+        """Validate that the patch got applied correctly."""
+        file = self.extraction_dir / "site-packages/numpy/file_added_via_patch.txt"
+        self.assertEqual(file.read_text(), "Hello from a patch!\n")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/pycross/py_wheel_library_test.py b/tests/pycross/py_wheel_library_test.py
index fa8e20e..25d896a 100644
--- a/tests/pycross/py_wheel_library_test.py
+++ b/tests/pycross/py_wheel_library_test.py
@@ -23,9 +23,7 @@ RUNFILES = runfiles.Create()
 class TestPyWheelLibrary(unittest.TestCase):
     def setUp(self):
         self.extraction_dir = Path(
-            RUNFILES.Rlocation(
-                "rules_python/tests/pycross/extracted_wheel_for_testing"
-            )
+            RUNFILES.Rlocation("rules_python/tests/pycross/extracted_wheel_for_testing")
         )
         self.assertTrue(self.extraction_dir.exists(), self.extraction_dir)
         self.assertTrue(self.extraction_dir.is_dir(), self.extraction_dir)
diff --git a/third_party/rules_pycross/pycross/private/tools/wheel_installer.py b/third_party/rules_pycross/pycross/private/tools/wheel_installer.py
index 8367f08..0c352cf 100644
--- a/third_party/rules_pycross/pycross/private/tools/wheel_installer.py
+++ b/third_party/rules_pycross/pycross/private/tools/wheel_installer.py
@@ -20,6 +20,7 @@ A tool that invokes pypa/build to build the given sdist tarball.
 import argparse
 import os
 import shutil
+import subprocess
 import sys
 import tempfile
 from pathlib import Path
@@ -97,6 +98,29 @@ def main(args: Any) -> None:
 
     setup_namespace_pkg_compatibility(lib_dir)
 
+    if args.patch:
+        if not args.patch_tool and not args.patch_tool_target:
+            raise ValueError("Specify one of 'patch_tool' or 'patch_tool_target'.")
+
+        patch_args = [
+            args.patch_tool or Path.cwd() / args.patch_tool_target
+        ] + args.patch_arg
+        for patch in args.patch:
+            with patch.open("r") as stdin:
+                try:
+                    subprocess.run(
+                        patch_args,
+                        stdin=stdin,
+                        check=True,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.STDOUT,
+                        cwd=args.directory,
+                    )
+                except subprocess.CalledProcessError as error:
+                    print(f"Patch {patch} failed to apply:")
+                    print(error.stdout.decode("utf-8"))
+                    raise
+
 
 def parse_flags(argv) -> Any:
     parser = argparse.ArgumentParser(description="Extract a Python wheel.")
@@ -127,6 +151,40 @@ def parse_flags(argv) -> Any:
         help="The output path.",
     )
 
+    parser.add_argument(
+        "--patch",
+        type=Path,
+        default=[],
+        action="append",
+        help="A patch file to apply.",
+    )
+
+    parser.add_argument(
+        "--patch-arg",
+        type=str,
+        default=[],
+        action="append",
+        help="An argument for the patch tool when applying the patches.",
+    )
+
+    parser.add_argument(
+        "--patch-tool",
+        type=str,
+        help=(
+            "The tool from PATH to invoke when applying patches. "
+            "If set, --patch-tool-target is ignored."
+        ),
+    )
+
+    parser.add_argument(
+        "--patch-tool-target",
+        type=Path,
+        help=(
+            "The path to the tool to invoke when applying patches. "
+            "Ignored when --patch-tool is set."
+        ),
+    )
+
     return parser.parse_args(argv[1:])
 
 
diff --git a/third_party/rules_pycross/pycross/private/wheel_library.bzl b/third_party/rules_pycross/pycross/private/wheel_library.bzl
index 381511a..166e1d0 100644
--- a/third_party/rules_pycross/pycross/private/wheel_library.bzl
+++ b/third_party/rules_pycross/pycross/private/wheel_library.bzl
@@ -33,19 +33,31 @@ def _py_wheel_library_impl(ctx):
     args = ctx.actions.args().use_param_file("--flagfile=%s")
     args.add("--wheel", wheel_file)
     args.add("--directory", out.path)
+    args.add_all(ctx.files.patches, format_each = "--patch=%s")
+    args.add_all(ctx.attr.patch_args, format_each = "--patch-arg=%s")
+    args.add("--patch-tool", ctx.attr.patch_tool)
 
-    inputs = [wheel_file]
+    tools = []
+    inputs = [wheel_file] + ctx.files.patches
     if name_file:
         inputs.append(name_file)
         args.add("--wheel-name-file", name_file)
 
+    if ctx.attr.patch_tool_target:
+        args.add("--patch-tool-target", ctx.attr.patch_tool_target.files_to_run.executable)
+        tools.append(ctx.executable.patch_tool_target)
+
     if ctx.attr.enable_implicit_namespace_pkgs:
         args.add("--enable-implicit-namespace-pkgs")
 
+    # We apply patches in the same action as the extraction to minimize the
+    # number of times we cache the wheel contents. If we were to split this
+    # into 2 actions, then the wheel contents would be cached twice.
     ctx.actions.run(
         inputs = inputs,
         outputs = [out],
         executable = ctx.executable._tool,
+        tools = tools,
         arguments = [args],
         # Set environment variables to make generated .pyc files reproducible.
         env = {
@@ -119,6 +131,31 @@ and py_test targets must specify either `legacy_create_init=False` or the global
 This option is required to support some packages which cannot handle the conversion to pkg-util style.
             """,
         ),
+        "patch_args": attr.string_list(
+            default = ["-p0"],
+            doc =
+                "The arguments given to the patch tool. Defaults to -p0, " +
+                "however -p1 will usually be needed for patches generated by " +
+                "git. If multiple -p arguments are specified, the last one will take effect.",
+        ),
+        "patch_tool": attr.string(
+            doc = "The patch(1) utility from the host to use. " +
+                  "If set, overrides `patch_tool_target`. Please note that setting " +
+                  "this means that builds are not completely hermetic.",
+        ),
+        "patch_tool_target": attr.label(
+            executable = True,
+            cfg = "exec",
+            doc = "The label of the patch(1) utility to use. " +
+                  "Only used if `patch_tool` is not set.",
+        ),
+        "patches": attr.label_list(
+            allow_files = True,
+            default = [],
+            doc =
+                "A list of files that are to be applied as patches after " +
+                "extracting the archive. This will use the patch command line tool.",
+        ),
         "python_version": attr.string(
             doc = "The python version required for this wheel ('PY2' or 'PY3')",
             values = ["PY2", "PY3", ""],