fix(pypi): pass requirements without env markers to the whl_library (#2488)

aignas · rickeylev · web-flow · commit bda710c54731 · 2024-12-12T00:17:52.000Z
With this change the environment markers from the requirements.txt files no longer end up in the whl_library definitions. I am reusing a function that already is parsing each requirement line for `sha256` values and added logic to extract the `marker` at that point. This means that the change is also trivial to backport to the `WORKSPACE` and the logic in the extension becomes simpler and we don't rely only on integration tests. Expected changes to the users: * If they have vendored pip requirements in `WORKSPACE`, those will be reformatted and the env markers will be removed. * The `MODULE.bazel.lock` file will be likewise reformatted if users are not using `--experimental_index_url`. Also, the env markers will not be passed in the `requirement`. * `bazel query 'deps("@pypi//foo")'` should start working in more cases. Fixes #2450. --------- Co-authored-by: Richard Levasseur <richardlev@gmail.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -60,6 +60,11 @@ Unreleased changes template.
 ### Fixed
 * (py_wheel) Use the default shell environment when building wheels to allow
   toolchains that search PATH to be used for the wheel builder tool.
+* (pypi) The requirement argument parsed to `whl_library` will now not have env
+  marker information allowing `bazel query` to work in cases where the `whl` is
+  available for all of the platforms and the sdist can be built. This fix is
+  for both WORKSPACE and `bzlmod` setups.
+  Fixes [#2450](https://github.com/bazelbuild/rules_python/issues/2450).
 
 {#v0-0-0-added}
 ### Added
diff --git a/examples/pip_parse_vendored/requirements.bzl b/examples/pip_parse_vendored/requirements.bzl
@@ -33,11 +33,11 @@ all_data_requirements = [
 ]
 
 _packages = [
-    ("my_project_pip_deps_vendored_certifi", "certifi==2023.7.22     --hash=sha256:539cc1d13202e33ca466e88b2807e29f4c13049d6d87031a3c110744495cb082     --hash=sha256:92d6037539857d8206b8f6ae472e8b77db8058fec5937a1ef3f54304089edbb9"),
-    ("my_project_pip_deps_vendored_charset_normalizer", "charset-normalizer==2.1.1     --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845     --hash=sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f"),
-    ("my_project_pip_deps_vendored_idna", "idna==3.4     --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4     --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2"),
-    ("my_project_pip_deps_vendored_requests", "requests==2.28.1     --hash=sha256:7c5599b102feddaa661c826c56ab4fee28bfd17f5abca1ebbe3e7f19d7c97983     --hash=sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349"),
-    ("my_project_pip_deps_vendored_urllib3", "urllib3==1.26.13     --hash=sha256:47cc05d99aaa09c9e72ed5809b60e7ba354e64b59c9c173ac3018642d8bb41fc     --hash=sha256:c083dd0dce68dbfbe1129d5271cb90f9447dea7d52097c6e0126120c521ddea8"),
+    ("my_project_pip_deps_vendored_certifi", "certifi==2023.7.22 --hash=sha256:539cc1d13202e33ca466e88b2807e29f4c13049d6d87031a3c110744495cb082 --hash=sha256:92d6037539857d8206b8f6ae472e8b77db8058fec5937a1ef3f54304089edbb9"),
+    ("my_project_pip_deps_vendored_charset_normalizer", "charset-normalizer==2.1.1 --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845 --hash=sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f"),
+    ("my_project_pip_deps_vendored_idna", "idna==3.4 --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2"),
+    ("my_project_pip_deps_vendored_requests", "requests==2.28.1 --hash=sha256:7c5599b102feddaa661c826c56ab4fee28bfd17f5abca1ebbe3e7f19d7c97983 --hash=sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349"),
+    ("my_project_pip_deps_vendored_urllib3", "urllib3==1.26.13 --hash=sha256:47cc05d99aaa09c9e72ed5809b60e7ba354e64b59c9c173ac3018642d8bb41fc --hash=sha256:c083dd0dce68dbfbe1129d5271cb90f9447dea7d52097c6e0126120c521ddea8"),
 ]
 _config = {
     "download_only": False,
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
@@ -321,10 +321,10 @@ def _create_whl_repos(
         for requirement in requirements:
             is_exposed = is_exposed or requirement.is_exposed
             if get_index_urls:
-                logger.warn(lambda: "falling back to pip for installing the right file for {}".format(requirement.requirement_line))
+                logger.warn(lambda: "falling back to pip for installing the right file for {}".format(requirement.srcs.requirement_line))
 
             args = dict(whl_library_args)  # make a copy
-            args["requirement"] = requirement.requirement_line
+            args["requirement"] = requirement.srcs.requirement_line
             if requirement.extra_pip_args:
                 args["extra_pip_args"] = requirement.extra_pip_args
 
diff --git a/python/private/pypi/index_sources.bzl b/python/private/pypi/index_sources.bzl
@@ -26,28 +26,43 @@ def index_sources(line):
         line(str): The requirements.txt entry.
 
     Returns:
-        A struct with shas attribute containing a list of shas to download from pypi_index.
+        A struct with shas attribute containing:
+            * `shas` - list[str]; shas to download from pypi_index.
+            * `version` - str; version of the package.
+            * `marker` - str; the marker expression, as per PEP508 spec.
+            * `requirement` - str; a requirement line without the marker. This can
+                be given to `pip` to install a package.
     """
+    line = line.replace("\\", " ")
     head, _, maybe_hashes = line.partition(";")
     _, _, version = head.partition("==")
     version = version.partition(" ")[0].strip()
 
-    if "@" in head:
-        shas = []
-    else:
-        maybe_hashes = maybe_hashes or line
-        shas = [
-            sha.strip()
-            for sha in maybe_hashes.split("--hash=sha256:")[1:]
-        ]
+    marker, _, _ = maybe_hashes.partition("--hash=")
+    maybe_hashes = maybe_hashes or line
+    shas = [
+        sha.strip()
+        for sha in maybe_hashes.split("--hash=sha256:")[1:]
+    ]
 
+    marker = marker.strip()
     if head == line:
-        head = line.partition("--hash=")[0].strip()
+        requirement = line.partition("--hash=")[0].strip()
     else:
-        head = head + ";" + maybe_hashes.partition("--hash=")[0].strip()
+        requirement = head.strip()
+
+    requirement_line = "{} {}".format(
+        requirement,
+        " ".join(["--hash=sha256:{}".format(sha) for sha in shas]),
+    ).strip()
+    if "@" in head:
+        requirement = requirement_line
+        shas = []
 
     return struct(
-        requirement = line if not shas else head,
+        requirement = requirement,
+        requirement_line = requirement_line,
         version = version,
         shas = sorted(shas),
+        marker = marker,
     )
diff --git a/python/private/pypi/parse_requirements.bzl b/python/private/pypi/parse_requirements.bzl
@@ -74,16 +74,22 @@ def parse_requirements(
         logger: repo_utils.logger or None, a simple struct to log diagnostic messages.
 
     Returns:
-        A tuple where the first element a dict of dicts where the first key is
-        the normalized distribution name (with underscores) and the second key
-        is the requirement_line, then value and the keys are structs with the
-        following attributes:
-         * distribution: The non-normalized distribution name.
-         * srcs: The Simple API downloadable source list.
-         * requirement_line: The original requirement line.
-         * target_platforms: The list of target platforms that this package is for.
-         * is_exposed: A boolean if the package should be exposed via the hub
+        {type}`dict[str, list[struct]]` where the key is the distribution name and the struct
+        contains the following attributes:
+         * `distribution`: {type}`str` The non-normalized distribution name.
+         * `srcs`: {type}`struct` The parsed requirement line for easier Simple
+           API downloading (see `index_sources` return value).
+         * `target_platforms`: {type}`list[str]` Target platforms that this package is for.
+             The format is `cp3{minor}_{os}_{arch}`.
+         * `is_exposed`: {type}`bool` `True` if the package should be exposed via the hub
            repository.
+         * `extra_pip_args`: {type}`list[str]` pip args to use in case we are
+           not using the bazel downloader to download the archives. This should
+           be passed to {obj}`whl_library`.
+         * `whls`: {type}`list[struct]` The list of whl entries that can be
+           downloaded using the bazel downloader.
+         * `sdist`: {type}`list[struct]` The sdist that can be downloaded using
+           the bazel downloader.
 
         The second element is extra_pip_args should be passed to `whl_library`.
     """
@@ -209,7 +215,6 @@ def parse_requirements(
                 struct(
                     distribution = r.distribution,
                     srcs = r.srcs,
-                    requirement_line = r.requirement_line,
                     target_platforms = sorted(target_platforms),
                     extra_pip_args = r.extra_pip_args,
                     whls = whls,
diff --git a/python/private/pypi/pip_repository.bzl b/python/private/pypi/pip_repository.bzl
@@ -101,7 +101,7 @@ def _pip_repository_impl(rctx):
         if not r:
             continue
         options = options or r.extra_pip_args
-        selected_requirements[name] = r.requirement_line
+        selected_requirements[name] = r.srcs.requirement_line
 
     bzl_packages = sorted(selected_requirements.keys())
 
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
@@ -28,7 +28,10 @@ def _mock_mctx(*modules, environ = {}, read = None):
             name = "unittest",
             arch = "exotic",
         ),
-        read = read or (lambda _: "simple==0.0.1 --hash=sha256:deadbeef --hash=sha256:deadbaaf"),
+        read = read or (lambda _: """\
+simple==0.0.1 \
+    --hash=sha256:deadbeef \
+    --hash=sha256:deadbaaf"""),
         modules = [
             struct(
                 name = modules[0].name,
@@ -262,7 +265,8 @@ def _test_simple_with_markers(env):
             read = lambda x: {
                 "universal.txt": """\
 torch==2.4.1+cpu ; platform_machine == 'x86_64'
-torch==2.4.1 ; platform_machine != 'x86_64'
+torch==2.4.1 ; platform_machine != 'x86_64' \
+    --hash=sha256:deadbeef
 """,
             }[x],
         ),
@@ -313,13 +317,13 @@ torch==2.4.1 ; platform_machine != 'x86_64'
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "repo": "pypi_315",
-            "requirement": "torch==2.4.1 ; platform_machine != 'x86_64'",
+            "requirement": "torch==2.4.1 --hash=sha256:deadbeef",
         },
         "pypi_315_torch_linux_x86_64_osx_x86_64_windows_x86_64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "repo": "pypi_315",
-            "requirement": "torch==2.4.1+cpu ; platform_machine == 'x86_64'",
+            "requirement": "torch==2.4.1+cpu",
         },
     })
     pypi.whl_mods().contains_exactly({})
@@ -351,16 +355,19 @@ def _test_download_only_multiple(env):
 --implementation=cp
 --abi=cp315
 
-simple==0.0.1 --hash=sha256:deadbeef
-extra==0.0.1 --hash=sha256:deadb00f
+simple==0.0.1 \
+    --hash=sha256:deadbeef
+extra==0.0.1 \
+    --hash=sha256:deadb00f
 """,
                 "requirements.osx_aarch64.txt": """\
 --platform=macosx_10_9_arm64
 --python-version=315
 --implementation=cp
 --abi=cp315
 
-simple==0.0.3 --hash=sha256:deadbaaf
+simple==0.0.3 \
+    --hash=sha256:deadbaaf
 """,
             }[x],
         ),
@@ -473,7 +480,9 @@ def _test_simple_get_index(env):
             ),
             read = lambda x: {
                 "requirements.txt": """
-simple==0.0.1 --hash=sha256:deadbeef --hash=sha256:deadb00f
+simple==0.0.1 \
+    --hash=sha256:deadbeef \
+    --hash=sha256:deadb00f
 some_pkg==0.0.1
 """,
             }[x],
diff --git a/tests/pypi/index_sources/index_sources_tests.bzl b/tests/pypi/index_sources/index_sources_tests.bzl
@@ -20,34 +20,62 @@ load("//python/private/pypi:index_sources.bzl", "index_sources")  # buildifier:
 _tests = []
 
 def _test_no_simple_api_sources(env):
-    inputs = [
-        "foo==0.0.1",
-        "foo==0.0.1 @ https://someurl.org",
-        "foo==0.0.1 @ https://someurl.org --hash=sha256:deadbeef",
-        "foo==0.0.1 @ https://someurl.org; python_version < 2.7 --hash=sha256:deadbeef",
-    ]
-    for input in inputs:
+    inputs = {
+        "foo==0.0.1": struct(
+            requirement = "foo==0.0.1",
+            marker = "",
+        ),
+        "foo==0.0.1 @ https://someurl.org": struct(
+            requirement = "foo==0.0.1 @ https://someurl.org",
+            marker = "",
+        ),
+        "foo==0.0.1 @ https://someurl.org --hash=sha256:deadbeef": struct(
+            requirement = "foo==0.0.1 @ https://someurl.org --hash=sha256:deadbeef",
+            marker = "",
+        ),
+        "foo==0.0.1 @ https://someurl.org; python_version < \"2.7\"\\    --hash=sha256:deadbeef": struct(
+            requirement = "foo==0.0.1 @ https://someurl.org --hash=sha256:deadbeef",
+            marker = "python_version < \"2.7\"",
+        ),
+    }
+    for input, want in inputs.items():
         got = index_sources(input)
         env.expect.that_collection(got.shas).contains_exactly([])
         env.expect.that_str(got.version).equals("0.0.1")
+        env.expect.that_str(got.requirement).equals(want.requirement)
+        env.expect.that_str(got.requirement_line).equals(got.requirement)
+        env.expect.that_str(got.marker).equals(want.marker)
 
 _tests.append(_test_no_simple_api_sources)
 
 def _test_simple_api_sources(env):
     tests = {
-        "foo==0.0.2 --hash=sha256:deafbeef    --hash=sha256:deadbeef": [
-            "deadbeef",
-            "deafbeef",
-        ],
-        "foo[extra]==0.0.2; (python_version < 2.7 or something_else == \"@\") --hash=sha256:deafbeef    --hash=sha256:deadbeef": [
-            "deadbeef",
-            "deafbeef",
-        ],
+        "foo==0.0.2 --hash=sha256:deafbeef    --hash=sha256:deadbeef": struct(
+            shas = [
+                "deadbeef",
+                "deafbeef",
+            ],
+            marker = "",
+            requirement = "foo==0.0.2",
+            requirement_line = "foo==0.0.2 --hash=sha256:deafbeef --hash=sha256:deadbeef",
+        ),
+        "foo[extra]==0.0.2; (python_version < 2.7 or extra == \"@\") --hash=sha256:deafbeef    --hash=sha256:deadbeef": struct(
+            shas = [
+                "deadbeef",
+                "deafbeef",
+            ],
+            marker = "(python_version < 2.7 or extra == \"@\")",
+            requirement = "foo[extra]==0.0.2",
+            requirement_line = "foo[extra]==0.0.2 --hash=sha256:deafbeef --hash=sha256:deadbeef",
+        ),
     }
-    for input, want_shas in tests.items():
+    for input, want in tests.items():
         got = index_sources(input)
-        env.expect.that_collection(got.shas).contains_exactly(want_shas)
+        env.expect.that_collection(got.shas).contains_exactly(want.shas)
         env.expect.that_str(got.version).equals("0.0.2")
+        env.expect.that_str(got.requirement).equals(want.requirement)
+        env.expect.that_str(got.requirement_line).equals(want.requirement_line)
+        env.expect.that_str(got.marker).equals(want.marker)
 
 _tests.append(_test_simple_api_sources)
 
diff --git a/tests/pypi/parse_requirements/parse_requirements_tests.bzl b/tests/pypi/parse_requirements/parse_requirements_tests.bzl
