From ff01f7e2903a85849f3278c20a0c14dba8004771 Mon Sep 17 00:00:00 2001
From: Josh Wilson <person142@users.noreply.github.com>
Date: Tue, 10 Nov 2020 20:31:50 -0800
Subject: [PATCH 1/2] Add `all_whl_requirements` to match `all_requirements`

---
 python/pip_install/README.md                    | 14 +++++++++++++-
 python/pip_install/extract_wheels/lib/BUILD     | 11 +++++++++++
 python/pip_install/extract_wheels/lib/bazel.py  | 11 ++++++++++-
 .../extract_wheels/lib/requirements_bzl_test.py | 17 +++++++++++++++++
 4 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 python/pip_install/extract_wheels/lib/requirements_bzl_test.py

diff --git a/python/pip_install/README.md b/python/pip_install/README.md
index 4db18529f8..cfa3b4fab9 100644
--- 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
index 4493bd1422..de67b2960e 100644
--- 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
index 964d4f91a7..0dabb6df13 100644
--- a/python/pip_install/extract_wheels/lib/bazel.py
+++ b/python/pip_install/extract_wheels/lib/bazel.py
@@ -74,6 +74,11 @@ 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}]
@@ -82,10 +87,14 @@ def requirement(name):
            name_key = name.replace("-", "_").replace(".", "_").lower()
            return "{repo}//pypi__" + name_key
 
+        all_whl_requirements = [{whl_requirement_labels}]
+
         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
new file mode 100644
index 0000000000..3424f3e9b7
--- /dev/null
+++ 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()

From b6886f26b7598ebe1e880230f36500b4c0f5a7b4 Mon Sep 17 00:00:00 2001
From: Josh Wilson <person142@users.noreply.github.com>
Date: Thu, 24 Dec 2020 18:45:49 -0800
Subject: [PATCH 2/2] Move all_whl_requirements to below all_requirements

---
 python/pip_install/extract_wheels/lib/bazel.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/pip_install/extract_wheels/lib/bazel.py b/python/pip_install/extract_wheels/lib/bazel.py
index 0dabb6df13..ef0d6e81b0 100644
--- a/python/pip_install/extract_wheels/lib/bazel.py
+++ b/python/pip_install/extract_wheels/lib/bazel.py
@@ -83,12 +83,12 @@ def generate_requirements_file_contents(repo_name: str, targets: Iterable[str])
         """\
         all_requirements = [{requirement_labels}]
 
+        all_whl_requirements = [{whl_requirement_labels}]
+
         def requirement(name):
            name_key = name.replace("-", "_").replace(".", "_").lower()
            return "{repo}//pypi__" + name_key
 
-        all_whl_requirements = [{whl_requirement_labels}]
-
         def whl_requirement(name):
             return requirement(name) + ":whl"
         """.format(