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

From edb6d0447edf6b6116c8aff73cde6117abbaf706 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         | 34 +++++++++++
 3 files changed, 109 insertions(+), 29 deletions(-)

diff --git a/python/pip_install/extract_wheels/extract_single_wheel.py b/python/pip_install/extract_wheels/extract_single_wheel.py
index a7cc672..26d4368 100644
--- a/python/pip_install/extract_wheels/extract_single_wheel.py
+++ b/python/pip_install/extract_wheels/extract_single_wheel.py
@@ -28,41 +28,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 5762cf5..46a94f5 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
@@ -81,7 +81,7 @@ def parse_whl_library_args(args: argparse.Namespace) -> Dict[str, Any]:
     whl_library_args.setdefault("python_interpreter", sys.executable)
 
     # These arguments are not used by `whl_library`
-    for arg in ("requirements_lock", "requirements_lock_label", "annotations"):
+    for arg in ("requirements_lock", "requirements_lock_label", "annotations", "overrides", "require_overrides"):
         if arg in whl_library_args:
             whl_library_args.pop(arg)
 
@@ -93,6 +93,8 @@ def generate_parsed_requirements_contents(
     repo_prefix: str,
     whl_library_args: Dict[str, Any],
     annotations: Dict[str, str] = dict(),
+    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
@@ -131,6 +133,13 @@ def generate_parsed_requirements_contents(
         _packages = {repo_names_and_reqs}
         _config = {args}
         _annotations = {annotations}
+        _overrides = {overrides}
+        _require_overrides = {require_overrides}
+
+        _NOP_OVERRIDE = {{
+            "url": None,
+            "sha256": None,
+        }}
 
         def _clean_name(name):
             return name.replace("-", "_").replace(".", "_").lower()
@@ -160,10 +169,22 @@ def generate_parsed_requirements_contents(
 
         def install_deps():
             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"],
                     **_config
                 )
         """.format(
@@ -178,6 +199,8 @@ def generate_parsed_requirements_contents(
             repo_names_and_reqs=repo_names_and_reqs,
             repo_prefix=repo_prefix,
             wheel_file_label=bazel.WHEEL_FILE_LABEL,
+            overrides=overrides or {},
+            require_overrides=require_overrides,
         )
     )
 
@@ -234,6 +257,16 @@ If set, it will take precedence over python_interpreter.",
         type=annotation.annotations_map_from_str_path,
         help="A json encoded file containing annotations for rendered packages.",
     )
+    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()
 
@@ -259,6 +292,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(
             """\
@@ -278,6 +316,8 @@ If set, it will take precedence over python_interpreter.",
             repo_prefix=args.repo_prefix,
             whl_library_args=whl_library_args,
             annotations=annotated_requirements,
+            overrides=overrides,
+            require_overrides=args.require_overrides,
         )
     )
 
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index d729ae9..afe3102 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -257,6 +257,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"
     else:
         args = [
@@ -391,6 +396,14 @@ pip_repository_attrs = {
         default = False,
         doc = "Create the repository in incremental mode.",
     ),
+    "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.",
@@ -483,6 +496,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)
 
@@ -515,6 +538,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)