Support annotations on pip packages with extras. (#865)

* Support annotations on pip packages with extras.

E.g., the following requirement:

```
requests[security]>=2.8.1
```

This is handled correctly by all of the other plumbing, but trying
to add an annotation to `requests` will fail. This is because
annotations have separate logic for parsing requirements in the
generated .bzl file. It would previously turn the requirement into
`requests[security]` rather than just `requests`.

* Add test verifying that annotations work for packages with extras.

examples/pip_parse_vendored/requirements.bzl

@@ -38,7 +38,7 @@ def entry_point(pkg, script = None):
 def _get_annotation(requirement):
     # This expects to parse `setuptools==58.2.0     --hash=sha256:2551203ae6955b9876741a26ab3e767bb3242dafe86a32a749ea0d78b6792f11`
     # down wo `setuptools`.
-    name = requirement.split(" ")[0].split("=")[0]
+    name = requirement.split(" ")[0].split("=")[0].split("[")[0]
     return _annotations.get(name)
 
 def install_deps(**whl_library_kwargs):

examples/pip_repository_annotations/.gitignore

@@ -0,0 +1 @@
+/bazel-*

examples/pip_repository_annotations/BUILD

@@ -16,9 +16,13 @@ compile_pip_requirements(
 py_test(
     name = "pip_parse_annotations_test",
     srcs = ["pip_repository_annotations_test.py"],
-    env = {"WHEEL_PKG_DIR": "pip_parsed_wheel"},
+    env = {
+        "REQUESTS_PKG_DIR": "pip_parsed_requests",
+        "WHEEL_PKG_DIR": "pip_parsed_wheel",
+    },
     main = "pip_repository_annotations_test.py",
     deps = [
+        "@pip_parsed_requests//:pkg",
         "@pip_parsed_wheel//:pkg",
         "@rules_python//python/runfiles",
     ],
@@ -27,10 +31,14 @@ py_test(
 py_test(
     name = "pip_install_annotations_test",
     srcs = ["pip_repository_annotations_test.py"],
-    env = {"WHEEL_PKG_DIR": "pip_installed_wheel"},
+    env = {
+        "REQUESTS_PKG_DIR": "pip_installed_requests",
+        "WHEEL_PKG_DIR": "pip_installed_wheel",
+    },
     main = "pip_repository_annotations_test.py",
     deps = [
         requirement("wheel"),
+        requirement("requests"),
         "@rules_python//python/runfiles",
     ],
 )

examples/pip_repository_annotations/WORKSPACE

@@ -30,6 +30,19 @@ load("@rules_python//python:pip.bzl", "package_annotation", "pip_install", "pip_
 # package. For details on `package_annotation` and it's uses, see the
 # docs at @rules_python//docs:pip.md`.
 ANNOTATIONS = {
+    # This annotation verifies that annotations work correctly for pip packages with extras
+    # specified, in this case requests[security].
+    "requests": package_annotation(
+        additive_build_content = """\
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
+write_file(
+    name = "generated_file",
+    out = "generated_file.txt",
+    content = ["Hello world from requests"],
+)
+""",
+        data = [":generated_file"],
+    ),
     "wheel": package_annotation(
         additive_build_content = """\
 load("@bazel_skylib//rules:write_file.bzl", "write_file")

examples/pip_repository_annotations/pip_repository_annotations_test.py

@@ -90,6 +90,26 @@ def test_data_exclude_glob(self):
             self.assertTrue(Path(metadata_path).exists())
             self.assertFalse(Path(wheel_path).exists())
 
+    def requests_pkg_dir(self) -> str:
+        env = os.environ.get("REQUESTS_PKG_DIR")
+        self.assertIsNotNone(env)
+        return env
+
+    def test_extra(self):
+        # This test verifies that annotations work correctly for pip packages with extras
+        # specified, in this case requests[security].
+        r = runfiles.Create()
+        rpath = r.Rlocation(
+            "pip_repository_annotations_example/external/{}/generated_file.txt".format(
+                self.requests_pkg_dir()
+            )
+        )
+        generated_file = Path(rpath)
+        self.assertTrue(generated_file.exists())
+
+        content = generated_file.read_text().rstrip()
+        self.assertEqual(content, "Hello world from requests")
+
 
 if __name__ == "__main__":
     unittest.main()

examples/pip_repository_annotations/requirements.in

@@ -3,3 +3,4 @@
 --extra-index-url https://pypi.python.org/simple/
 
 wheel
+requests[security]>=2.8.1

examples/pip_repository_annotations/requirements.txt

@@ -6,6 +6,26 @@
 #
 --extra-index-url https://pypi.python.org/simple/
 
+certifi==2022.9.24 \
+    --hash=sha256:0d9c601124e5a6ba9712dbc60d9c53c21e34f5f641fe83002317394311bdce14 \
+    --hash=sha256:90c1a32f1d68f940488354e36370f6cca89f0f106db09518524c88d6ed83f382
+    # via requests
+charset-normalizer==2.1.1 \
+    --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845 \
+    --hash=sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f
+    # via requests
+idna==3.4 \
+    --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \
+    --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2
+    # via requests
+requests[security]==2.28.1 \
+    --hash=sha256:7c5599b102feddaa661c826c56ab4fee28bfd17f5abca1ebbe3e7f19d7c97983 \
+    --hash=sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349
+    # via -r ./requirements.in
+urllib3==1.26.12 \
+    --hash=sha256:3fa96cf423e6987997fc326ae8df396db2a8b7c667747d47ddd8ecba91f4a74e \
+    --hash=sha256:b930dd878d5a8afb066a637fbb35144fe7901e3b209d1cd4f524bd0e9deee997
+    # via requests
 wheel==0.37.1 \
     --hash=sha256:4bdcd7d840138086126cd09254dc6195fb4fc6f01c050a1d7236f2630db1d22a \
     --hash=sha256:e9a504e793efbca1b8e0e9cb979a249cf4a0a7b5b8c9e8b65a5e39d49529c1c4

python/pip_install/extract_wheels/parse_requirements_to_bzl.py

@@ -155,7 +155,7 @@ def entry_point(pkg, script = None):
         def _get_annotation(requirement):
             # This expects to parse `setuptools==58.2.0     --hash=sha256:2551203ae6955b9876741a26ab3e767bb3242dafe86a32a749ea0d78b6792f11`
             # down wo `setuptools`.
-            name = requirement.split(" ")[0].split("=")[0]
+            name = requirement.split(" ")[0].split("=")[0].split("[")[0]
             return _annotations.get(name)
 
         def install_deps(**whl_library_kwargs):