third_party/rules_python/0001-Support-overriding-individual-packages.patch

From 662f59afaecd7ecff5bd5234c8bbd9c219b7f24f Mon Sep 17 00:00:00 2001
From: Philipp Schrader <philipp.schrader@gmail.com>
Date: Sun, 11 Sep 2022 22:04:47 -0700
Subject: [PATCH] Support overriding individual packages

---
 .../extract_wheels/extract_single_wheel.py    | 60 ++++++++++---------
 .../parse_requirements_to_bzl.py              | 44 +++++++++++++-
 python/pip_install/pip_repository.bzl         | 38 ++++++++++++
 3 files changed, 114 insertions(+), 28 deletions(-)

diff --git a/python/pip_install/extract_wheels/extract_single_wheel.py b/python/pip_install/extract_wheels/extract_single_wheel.py
index ff64291..8742d25 100644
--- a/python/pip_install/extract_wheels/extract_single_wheel.py
+++ b/python/pip_install/extract_wheels/extract_single_wheel.py
@@ -50,41 +50,47 @@ def main() -> None:
         type=annotation_from_str_path,
         help="A json encoded file containing annotations for rendered packages.",
     )
+    parser.add_argument(
+        "--pre-downloaded",
+        action="store_true",
+        help="If set, skips the pip download step. The .whl file is assumbed to be downloaded by bazel.",
+    )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
     deserialized_args = dict(vars(args))
     arguments.deserialize_structured_args(deserialized_args)
 
-    configure_reproducible_wheels()
+    if not args.pre_downloaded:
+        configure_reproducible_wheels()
 
-    pip_args = (
-        [sys.executable, "-m", "pip"]
-        + (["--isolated"] if args.isolated else [])
-        + ["download" if args.download_only else "wheel", "--no-deps"]
-        + deserialized_args["extra_pip_args"]
-    )
+        pip_args = (
+            [sys.executable, "-m", "pip"]
+            + (["--isolated"] if args.isolated else [])
+            + ["download" if args.download_only else "wheel", "--no-deps"]
+            + deserialized_args["extra_pip_args"]
+        )
 
-    requirement_file = NamedTemporaryFile(mode="wb", delete=False)
-    try:
-        requirement_file.write(args.requirement.encode("utf-8"))
-        requirement_file.flush()
-        # Close the file so pip is allowed to read it when running on Windows.
-        # For more information, see: https://bugs.python.org/issue14243
-        requirement_file.close()
-        # Requirement specific args like --hash can only be passed in a requirements file,
-        # so write our single requirement into a temp file in case it has any of those flags.
-        pip_args.extend(["-r", requirement_file.name])
-
-        env = os.environ.copy()
-        env.update(deserialized_args["environment"])
-        # Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
-        subprocess.run(pip_args, check=True, env=env)
-    finally:
+        requirement_file = NamedTemporaryFile(mode="wb", delete=False)
         try:
-            os.unlink(requirement_file.name)
-        except OSError as e:
-            if e.errno != errno.ENOENT:
-                raise
+            requirement_file.write(args.requirement.encode("utf-8"))
+            requirement_file.flush()
+            # Close the file so pip is allowed to read it when running on Windows.
+            # For more information, see: https://bugs.python.org/issue14243
+            requirement_file.close()
+            # Requirement specific args like --hash can only be passed in a requirements file,
+            # so write our single requirement into a temp file in case it has any of those flags.
+            pip_args.extend(["-r", requirement_file.name])
+
+            env = os.environ.copy()
+            env.update(deserialized_args["environment"])
+            # Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
+            subprocess.run(pip_args, check=True, env=env)
+        finally:
+            try:
+                os.unlink(requirement_file.name)
+            except OSError as e:
+                if e.errno != errno.ENOENT:
+                    raise
 
     name, extras_for_pkg = requirements._parse_requirement_for_extra(args.requirement)
     extras = {name: extras_for_pkg} if extras_for_pkg and name else dict()
diff --git a/python/pip_install/extract_wheels/parse_requirements_to_bzl.py b/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
index 686a57d..60936a9 100644
--- a/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
+++ b/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
@@ -4,7 +4,7 @@ import shlex
 import sys
 import textwrap
 from pathlib import Path
-from typing import Any, Dict, List, TextIO, Tuple
+from typing import Any, Dict, List, Optional, TextIO, Tuple
 
 from pip._internal.network.session import PipSession
 from pip._internal.req import constructors
@@ -86,6 +86,8 @@ def parse_whl_library_args(args: argparse.Namespace) -> Dict[str, Any]:
         "requirements_lock_label",
         "annotations",
         "bzlmod",
+        "overrides",
+        "require_overrides",
     ):
         if arg in whl_library_args:
             whl_library_args.pop(arg)
@@ -100,6 +102,8 @@ def generate_parsed_requirements_contents(
     whl_library_args: Dict[str, Any],
     annotations: Dict[str, str] = dict(),
     bzlmod: bool = False,
+    overrides: Optional[Dict[str, Dict[str, str]]] = None,
+    require_overrides: bool = False,
 ) -> str:
     """
     Parse each requirement from the requirements_lock file, and prepare arguments for each
@@ -133,10 +137,22 @@ def generate_parsed_requirements_contents(
             whl_config = dict(_config)
             whl_config.update(whl_library_kwargs)
             for name, requirement in _packages:
+                override_entry = requirement.split(" ")[0]
+                override_name, _, version = override_entry.partition("==")
+                override_key = "%s==%s" % (_clean_name(override_name), version)
+                override = _overrides.get(override_key)
+                if not override:
+                    if _require_overrides:
+                        fail("Failed to find an override for \\"{{}}\\" in the \\"overrides\\" JSON file".format(override_key))
+                    else:
+                        override = _NOP_OVERRIDE
+
                 whl_library(
                     name = name,
                     requirement = requirement,
                     annotation = _get_annotation(requirement),
+                    url = override["url"],
+                    sha256 = override["sha256"],
                     **whl_config
                 )
 """
@@ -154,6 +170,13 @@ def generate_parsed_requirements_contents(
         _config = {args}
         _annotations = {annotations}
         _bzlmod = {bzlmod}
+        _overrides = {overrides}
+        _require_overrides = {require_overrides}
+
+        _NOP_OVERRIDE = {{
+            "url": None,
+            "sha256": None,
+        }}
 
         def _clean_name(name):
             return name.replace("-", "_").replace(".", "_").lower()
@@ -204,6 +227,8 @@ def generate_parsed_requirements_contents(
             repo_prefix=repo_prefix,
             wheel_file_label=bazel.WHEEL_FILE_LABEL,
             bzlmod=bzlmod,
+            overrides=overrides or {},
+            require_overrides=require_overrides,
         )
     )
 
@@ -266,6 +291,16 @@ If set, it will take precedence over python_interpreter.",
         default=False,
         help="Whether this script is run under bzlmod. Under bzlmod we don't generate the install_deps() macro as it isn't needed.",
     )
+    parser.add_argument(
+        "--overrides",
+        type=Path,
+        help="A json encoded file containing URL overrides for packages.",
+    )
+    parser.add_argument(
+        "--require-overrides",
+        action="store_true",
+        help="If set, requires that every requirement has a URL override in the --overrides JSON file.",
+    )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
 
@@ -291,6 +326,11 @@ If set, it will take precedence over python_interpreter.",
             }
         )
 
+    if args.overrides:
+        overrides = json.loads(args.overrides.read_text())
+    else:
+        overrides = None
+
     output.write(
         textwrap.dedent(
             """\
@@ -313,6 +353,8 @@ If set, it will take precedence over python_interpreter.",
             whl_library_args=whl_library_args,
             annotations=annotated_requirements,
             bzlmod=args.bzlmod,
+            overrides=overrides,
+            require_overrides=args.require_overrides,
         )
     )
 
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index 7fbf503..5af0731 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -322,6 +322,11 @@ def _pip_repository_impl(rctx):
     args += ["--python_interpreter", _get_python_interpreter_attr(rctx)]
     if rctx.attr.python_interpreter_target:
         args += ["--python_interpreter_target", str(rctx.attr.python_interpreter_target)]
+    if rctx.attr.overrides:
+        overrides_file = rctx.path(rctx.attr.overrides).realpath
+        args += ["--overrides", overrides_file]
+    if rctx.attr.require_overrides:
+        args += ["--require-overrides"]
     progress_message = "Parsing requirements to starlark"
 
     args += ["--repo", rctx.attr.name, "--repo-prefix", rctx.attr.repo_prefix]
@@ -447,6 +452,18 @@ pip_repository_attrs = {
 we do not create the install_deps() macro.
 """,
     ),
+    "overrides": attr.label(
+        allow_single_file = True,
+        doc = "A JSON file containing overrides. TBD",
+    ),
+    "require_overrides": attr.bool(
+        default = False,
+        doc = "If True, every requirement must have an entry in the \"overrides\" JSON file.",
+    ),
+    "requirements": attr.label(
+        allow_single_file = True,
+        doc = "A 'requirements.txt' pip requirements file.",
+    ),
     "requirements_darwin": attr.label(
         allow_single_file = True,
         doc = "Override the requirements_lock attribute when the host platform is Mac OS",
@@ -535,6 +552,16 @@ def _whl_library_impl(rctx):
             "--annotation",
             rctx.path(rctx.attr.annotation),
         ])
+    if rctx.attr.url:
+        basename = rctx.attr.url.split("/")[-1]
+        download_result = rctx.download(
+            output = basename,
+            url = rctx.attr.url,
+            sha256 = rctx.attr.sha256 or None,
+        )
+        if not download_result.success:
+            fail("Failed to download {}".format(rctx.attr.url))
+        args.append("--pre-downloaded")
 
     args = _parse_optional_attrs(rctx, args)
 
@@ -567,6 +594,17 @@ whl_library_attrs = {
         mandatory = True,
         doc = "Python requirement string describing the package to make available",
     ),
+    "url": attr.string(
+        doc = (
+            "Set this to download the package from the specified URL instead of using pip. "
+        ),
+    ),
+    "sha256": attr.string(
+        doc  = (
+            "Optionally set this when using the 'url' parameter. " +
+            "Must be the expected checksum of the downloaded file."
+        ),
+    )
 }
 
 whl_library_attrs.update(**common_attrs)