Add all_whl_requirements to match all_requirements (bazel-contrib#377)

person142 · web-flow · commit c97265590fd1 · 2021-01-06T09:55:35.000+11:00
diff --git a/python/pip_install/README.md b/python/pip_install/README.md
@@ -64,7 +64,12 @@ pip_install(
 #### Example `BUILD` file.
 
 ```python
-load("@py_deps//:requirements.bzl", "requirement", "whl_requirement")
+load(
+    "@py_deps//:requirements.bzl",
+    "requirement",
+    "whl_requirement",
+    "all_whl_requirements",
+)
 
 py_binary(
     name = "main",
@@ -82,6 +87,13 @@ filegroup(
         whl_requirement("boto3"),
     ]
 )
+
+# If you need all of the wheels, say to upload them to your own
+# private wheelhouse, you can use all_whl_requirements.
+filegroup(
+    name = "all_whls",
+    data = all_whl_requirements,
+)
 ```
 
 Note that above you do not need to add transitively required packages to `deps = [ ... ]` or `data = [ ... ]`
diff --git a/python/pip_install/extract_wheels/lib/BUILD b/python/pip_install/extract_wheels/lib/BUILD
@@ -54,6 +54,17 @@ py_test(
     data = ["//experimental/examples/wheel:minimal_with_py_package"]
 )
 
+py_test(
+    name = "requirements_bzl_test",
+    size = "small",
+    srcs = [
+        "requirements_bzl_test.py",
+    ],
+    deps = [
+        ":lib",
+    ],
+)
+
 filegroup(
     name = "distribution",
     srcs = glob(
diff --git a/python/pip_install/extract_wheels/lib/bazel.py b/python/pip_install/extract_wheels/lib/bazel.py
@@ -74,18 +74,27 @@ def generate_requirements_file_contents(repo_name: str, targets: Iterable[str])
         A complete requirements.bzl file as a string
     """
 
+    sorted_targets = sorted(targets)
+    requirement_labels = ",".join(sorted_targets)
+    whl_requirement_labels = ",".join(
+        '"{}:whl"'.format(target.strip('"')) for target in sorted_targets
+    )
     return textwrap.dedent(
         """\
         all_requirements = [{requirement_labels}]
 
+        all_whl_requirements = [{whl_requirement_labels}]
+
         def requirement(name):
            name_key = name.replace("-", "_").replace(".", "_").lower()
            return "{repo}//pypi__" + name_key
 
         def whl_requirement(name):
             return requirement(name) + ":whl"
         """.format(
-            repo=repo_name, requirement_labels=",".join(sorted(targets))
+            repo=repo_name,
+            requirement_labels=requirement_labels,
+            whl_requirement_labels=whl_requirement_labels,
         )
     )
 
diff --git a/python/pip_install/extract_wheels/lib/requirements_bzl_test.py b/python/pip_install/extract_wheels/lib/requirements_bzl_test.py
@@ -0,0 +1,17 @@
+import unittest
+
+from python.pip_install.extract_wheels.lib import bazel
+
+
+class TestGenerateRequirementsFileContents(unittest.TestCase):
+    def test_all_wheel_requirements(self) -> None:
+        contents = bazel.generate_requirements_file_contents(
+            repo_name='test',
+            targets=['"@test//pypi__pkg1"', '"@test//pypi__pkg2"'],
+        )
+        expected = 'all_whl_requirements = ["@test//pypi__pkg1:whl","@test//pypi__pkg2:whl"]'
+        self.assertIn(expected, contents)
+
+
+if __name__ == "__main__":
+    unittest.main()
