Fix for requirements_lock with PEP440 direct references (bazel-contrib#756)

groodt · web-flow · commit 834149dfcd9e · 2022-07-15T06:09:21.000+10:00
diff --git a/python/pip_install/extract_wheels/parse_requirements_to_bzl.py b/python/pip_install/extract_wheels/parse_requirements_to_bzl.py
@@ -37,7 +37,12 @@ def parse_install_requirements(
     ):
         if parsed_line.is_requirement:
             install_req = constructors.install_req_from_line(parsed_line.requirement)
-            if not install_req.is_pinned:
+            if (
+                # PEP-440 direct references are considered pinned
+                # See: https://peps.python.org/pep-0440/#direct-references and https://peps.python.org/pep-0508/
+                not install_req.link and
+                not install_req.is_pinned
+            ):
                 unpinned_reqs.append(str(install_req))
             install_req_and_lines.append(
                 (install_req, line)
diff --git a/python/pip_install/extract_wheels/parse_requirements_to_bzl_test.py b/python/pip_install/extract_wheels/parse_requirements_to_bzl_test.py
@@ -119,6 +119,33 @@ def test_parse_install_requirements_with_args(self):
                     ),
                 )
 
+    def test_parse_install_requirements_pinned_direct_reference(self):
+        # Test PEP-440 direct references
+        with tempfile.TemporaryDirectory() as temp_dir:
+            requirements_lock = Path(temp_dir) / "requirements.txt"
+            requirements_lock.write_text(
+                dedent(
+                    """\
+                onnx @ https://files.pythonhosted.org/packages/24/93/f5b001dc0f5de84ce049a34ff382032cd9478e1080aa6ac48470fa810577/onnx-1.11.0-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl \
+                    --hash=sha256:67c6d2654c1c203e5c839a47900b51f588fd0de71bbd497fb193d30a0b3ec1e9
+                """
+                )
+            )
+
+            install_req_and_lines = parse_install_requirements(
+                str(requirements_lock), ["-v"]
+            )
+
+            self.assertEqual(len(install_req_and_lines), 1)
+            self.assertEqual(install_req_and_lines[0][0].name, "onnx")
+
+            self.assertTupleEqual(
+                install_req_and_lines[0][1:],
+                (
+                    "onnx @ https://files.pythonhosted.org/packages/24/93/f5b001dc0f5de84ce049a34ff382032cd9478e1080aa6ac48470fa810577/onnx-1.11.0-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl                     --hash=sha256:67c6d2654c1c203e5c839a47900b51f588fd0de71bbd497fb193d30a0b3ec1e9",
+                ),
+            )
+
 
 if __name__ == "__main__":
     unittest.main()
