python: Fix issues with top-level deps that have extras
I824d9ec51c562f0a06b8d2ec1fc27a0e1c8a2ad9 imports the `jax[cuda12]`
module. That's currently encountering this error:
$ bazel build --repo_env=FRC971_RUNNING_IN_CI=1 @pip_deps_jax//...
...
ERROR: Traceback (most recent call last):
File "/bazel-cache/buildkite-agent/builds/frc8971-1/spartan-robotics/971-robot-code/WORKSPACE", line 96, column 17, in <toplevel>
install_pip_deps()
File "/bazel-cache/buildkite-agent/builds/frc8971-1/spartan-robotics/k8_output_base/external/pip_deps/requirements.bzl", line 70, column 21, in install_deps
fail("Failed to find an override for \"{}\" in the \"overrides\" JSON file".format(override_key))
Error in fail: Failed to find an override for "jax[cuda12]==0.4.30" in the "overrides" JSON file
The issue here is that overrides logic doesn't ignore extras. We don't
need extras information in this context. This patch fixes the issue by
removing the extras information from top-level requirements when
looking up URL overrides.
Signed-off-by: Philipp Schrader <philipp.schrader@gmail.com>
Change-Id: I1c841e0b1e640824ab23939544195a29f2486bce
diff --git a/third_party/rules_python/0001-Support-overriding-individual-packages.patch b/third_party/rules_python/0001-Support-overriding-individual-packages.patch
index af89ecc..6f01b1b 100644
--- a/third_party/rules_python/0001-Support-overriding-individual-packages.patch
+++ b/third_party/rules_python/0001-Support-overriding-individual-packages.patch
@@ -1,16 +1,16 @@
-From 662f59afaecd7ecff5bd5234c8bbd9c219b7f24f Mon Sep 17 00:00:00 2001
+From b9b8deb69a6c53a0d688988161cd057e6f94881b 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 +++++++++++++-
+ .../parse_requirements_to_bzl.py | 51 +++++++++++++++-
python/pip_install/pip_repository.bzl | 38 ++++++++++++
- 3 files changed, 114 insertions(+), 28 deletions(-)
+ 3 files changed, 121 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
+index ff642910..8742d250 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:
@@ -89,7 +89,7 @@
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
+index 686a57d8..002e6857 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
@@ -125,7 +125,7 @@
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_key = "%s==%s" % (_clean_extras(_clean_name(override_name)), version)
+ override = _overrides.get(override_key)
+ if not override:
+ if _require_overrides:
@@ -142,7 +142,7 @@
**whl_config
)
"""
-@@ -154,6 +170,13 @@ def generate_parsed_requirements_contents(
+@@ -154,10 +170,24 @@ def generate_parsed_requirements_contents(
_config = {args}
_annotations = {annotations}
_bzlmod = {bzlmod}
@@ -156,7 +156,18 @@
def _clean_name(name):
return name.replace("-", "_").replace(".", "_").lower()
-@@ -204,6 +227,8 @@ def generate_parsed_requirements_contents(
+
++ def _clean_extras(name):
++ bracket_start = name.find("[")
++ bracket_end = name.find("]")
++ if bracket_start == -1 or bracket_end == -1:
++ return name
++ return name[:bracket_start] + name[bracket_end + 1:]
++
+ def requirement(name):
+ if _bzlmod:
+ return "@@{repo}//:" + _clean_name(name) + "_{py_library_label}"
+@@ -204,6 +234,8 @@ def generate_parsed_requirements_contents(
repo_prefix=repo_prefix,
wheel_file_label=bazel.WHEEL_FILE_LABEL,
bzlmod=bzlmod,
@@ -165,7 +176,7 @@
)
)
-@@ -266,6 +291,16 @@ If set, it will take precedence over python_interpreter.",
+@@ -266,6 +298,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.",
)
@@ -182,7 +193,7 @@
arguments.parse_common_args(parser)
args = parser.parse_args()
-@@ -291,6 +326,11 @@ If set, it will take precedence over python_interpreter.",
+@@ -291,6 +333,11 @@ If set, it will take precedence over python_interpreter.",
}
)
@@ -194,7 +205,7 @@
output.write(
textwrap.dedent(
"""\
-@@ -313,6 +353,8 @@ If set, it will take precedence over python_interpreter.",
+@@ -313,6 +360,8 @@ If set, it will take precedence over python_interpreter.",
whl_library_args=whl_library_args,
annotations=annotated_requirements,
bzlmod=args.bzlmod,
@@ -204,7 +215,7 @@
)
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
-index 7fbf503..5af0731 100644
+index 7fbf5039..5af07315 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -322,6 +322,11 @@ def _pip_repository_impl(rctx):
diff --git a/third_party/rules_python/0002-Allow-user-to-patch-wheels.patch b/third_party/rules_python/0002-Allow-user-to-patch-wheels.patch
index 9db2d9e..3669ca5 100644
--- a/third_party/rules_python/0002-Allow-user-to-patch-wheels.patch
+++ b/third_party/rules_python/0002-Allow-user-to-patch-wheels.patch
@@ -1,4 +1,4 @@
-From f828c9aad94b56655b352c4bed9b475d7430865e Mon Sep 17 00:00:00 2001
+From e9002bef44df84370649b995b8c9e6a89d4d37e9 Mon Sep 17 00:00:00 2001
From: Philipp Schrader <philipp.schrader@gmail.com>
Date: Sat, 24 Sep 2022 15:56:33 -0700
Subject: [PATCH] Allow user to patch wheels
@@ -25,7 +25,7 @@
5 files changed, 100 insertions(+), 3 deletions(-)
diff --git a/python/pip_install/extract_wheels/annotation.py b/python/pip_install/extract_wheels/annotation.py
-index 48aaa80..fe8b4dc 100644
+index 48aaa802..fe8b4dc5 100644
--- a/python/pip_install/extract_wheels/annotation.py
+++ b/python/pip_install/extract_wheels/annotation.py
@@ -19,6 +19,7 @@ class Annotation(OrderedDict):
@@ -48,7 +48,7 @@
class AnnotationsMap:
"""A mapping of python package names to [Annotation]"""
diff --git a/python/pip_install/extract_wheels/bazel.py b/python/pip_install/extract_wheels/bazel.py
-index 8f442c9..f4b7f26 100644
+index 8f442c93..f4b7f26a 100644
--- a/python/pip_install/extract_wheels/bazel.py
+++ b/python/pip_install/extract_wheels/bazel.py
@@ -2,6 +2,7 @@
@@ -107,7 +107,7 @@
os.remove(whl.path)
return f"//{directory}"
diff --git a/python/pip_install/extract_wheels/extract_single_wheel.py b/python/pip_install/extract_wheels/extract_single_wheel.py
-index 8742d25..50a1243 100644
+index 8742d250..50a1243e 100644
--- a/python/pip_install/extract_wheels/extract_single_wheel.py
+++ b/python/pip_install/extract_wheels/extract_single_wheel.py
@@ -1,4 +1,5 @@
@@ -166,7 +166,7 @@
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 60936a9..3dd179b 100644
+index 002e6857..fc7fe780 100644
--- a/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
+++ b/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
@@ -88,6 +88,7 @@ def parse_whl_library_args(args: argparse.Namespace) -> Dict[str, Any]:
@@ -214,7 +214,7 @@
_NOP_OVERRIDE = {{
"url": None,
-@@ -229,6 +238,7 @@ def generate_parsed_requirements_contents(
+@@ -236,6 +245,7 @@ def generate_parsed_requirements_contents(
bzlmod=bzlmod,
overrides=overrides or {},
require_overrides=require_overrides,
@@ -222,7 +222,7 @@
)
)
-@@ -301,6 +311,13 @@ If set, it will take precedence over python_interpreter.",
+@@ -308,6 +318,13 @@ If set, it will take precedence over python_interpreter.",
action="store_true",
help="If set, requires that every requirement has a URL override in the --overrides JSON file.",
)
@@ -236,7 +236,7 @@
arguments.parse_common_args(parser)
args = parser.parse_args()
-@@ -331,6 +348,11 @@ If set, it will take precedence over python_interpreter.",
+@@ -338,6 +355,11 @@ If set, it will take precedence over python_interpreter.",
else:
overrides = None
@@ -248,7 +248,7 @@
output.write(
textwrap.dedent(
"""\
-@@ -355,6 +377,7 @@ If set, it will take precedence over python_interpreter.",
+@@ -362,6 +384,7 @@ If set, it will take precedence over python_interpreter.",
bzlmod=args.bzlmod,
overrides=overrides,
require_overrides=args.require_overrides,
@@ -257,7 +257,7 @@
)
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
-index 5af0731..bf7f99a 100644
+index 5af07315..bf7f99a8 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -327,6 +327,9 @@ def _pip_repository_impl(rctx):