Add support for requirement extras (#28)

Author: dillon-giacoppo

extract_wheels/__init__.py

@@ -11,7 +11,7 @@
 import subprocess
 import sys
 
-from extract_wheels.lib import bazel
+from extract_wheels.lib import bazel, requirements
 
 
 def configure_reproducible_wheels() -> None:
@@ -70,8 +70,10 @@ def main() -> None:
         [sys.executable, "-m", "pip", "wheel", "-r", args.requirements]
     )
 
+    extras = requirements.parse_extras(args.requirements)
+
     targets = [
-        '"%s%s"' % (args.repo, bazel.extract_wheel(whl, []))
+        '"%s%s"' % (args.repo, bazel.extract_wheel(whl, extras))
         for whl in glob.glob("*.whl")
     ]
 

extract_wheels/lib/BUILD

@@ -7,6 +7,7 @@ py_library(
         "bazel.py",
         "namespace_pkgs.py",
         "purelib.py",
+        "requirements.py",
         "wheel.py",
     ],
     deps = [
@@ -26,3 +27,15 @@ py_test(
         ":lib",
     ],
 )
+
+py_test(
+    name = "requirements_test",
+    size = "small",
+    srcs = [
+        "requirements_test.py",
+    ],
+    tags = ["unit"],
+    deps = [
+        ":lib",
+    ],
+)

extract_wheels/lib/bazel.py

@@ -1,7 +1,7 @@
 """Utility functions to manipulate Bazel files"""
 import os
 import textwrap
-from typing import Iterable, List
+from typing import Iterable, List, Dict, Set
 
 from extract_wheels.lib import namespace_pkgs, wheel, purelib
 
@@ -113,7 +113,7 @@ def setup_namespace_pkg_compatibility(wheel_dir: str) -> None:
         namespace_pkgs.add_pkgutil_style_namespace_pkg_init(ns_pkg_dir)
 
 
-def extract_wheel(wheel_file: str, extras: List[str]) -> str:
+def extract_wheel(wheel_file: str, extras: Dict[str, Set[str]]) -> str:
     """Extracts wheel into given directory and creates a py_library target.
 
     Args:
@@ -134,16 +134,17 @@ def extract_wheel(wheel_file: str, extras: List[str]) -> str:
     purelib.spread_purelib_into_root(directory)
     setup_namespace_pkg_compatibility(directory)
 
+    extras_requested = extras[whl.name] if whl.name in extras else set()
+
+    sanitised_dependencies = [
+        '"//%s"' % sanitise_name(d) for d in sorted(whl.dependencies(extras_requested))
+    ]
+
     with open(os.path.join(directory, "BUILD"), "w") as build_file:
-        build_file.write(
-            generate_build_file_contents(
-                sanitise_name(whl.name),
-                [
-                    '"//%s"' % sanitise_name(d)
-                    for d in sorted(whl.dependencies(extras_requested=extras))
-                ],
-            )
+        contents = generate_build_file_contents(
+            sanitise_name(whl.name), sanitised_dependencies,
         )
+        build_file.write(contents)
 
     os.remove(whl.path)
 

extract_wheels/lib/namespace_pkgs_test.py

@@ -1,7 +1,5 @@
-import os
 import pathlib
 import shutil
-import sys
 import tempfile
 import unittest
 
@@ -133,17 +131,5 @@ def test_empty_case(self):
         self.assertEqual(actual, set())
 
 
-def main():
-    loader = unittest.TestLoader()
-    cur_dir = os.path.dirname(os.path.realpath(__file__))
-
-    suite = loader.discover(cur_dir)
-
-    runner = unittest.TextTestRunner()
-    result = runner.run(suite)
-    if result.errors or result.failures:
-        sys.exit(1)
-
-
 if __name__ == "__main__":
-    main()
+    unittest.main()

extract_wheels/lib/requirements.py

@@ -0,0 +1,45 @@
+import re
+from typing import Dict, Set, Tuple, Optional
+
+
+def parse_extras(requirements_path: str) -> Dict[str, Set[str]]:
+    """Parse over the requirements.txt file to find extras requested.
+
+    Args:
+        requirements_path: The filepath for the requirements.txt file to parse.
+
+    Returns:
+         A dictionary mapping the requirement name to a set of extras requested.
+    """
+
+    extras_requested = {}
+    with open(requirements_path, "r") as requirements:
+        # Merge all backslash line continuations so we parse each requirement as a single line.
+        for line in requirements.read().replace("\\\n", "").split("\n"):
+            requirement, extras = _parse_requirement_for_extra(line)
+            if requirement and extras:
+                extras_requested[requirement] = extras
+
+    return extras_requested
+
+
+def _parse_requirement_for_extra(
+    requirement: str,
+) -> Tuple[Optional[str], Optional[Set[str]]]:
+    """Given a requirement string, returns the requirement name and set of extras, if extras specified.
+    Else, returns (None, None)
+    """
+
+    # https://www.python.org/dev/peps/pep-0508/#grammar
+    extras_pattern = re.compile(
+        r"^\s*([0-9A-Za-z][0-9A-Za-z_.\-]*)\s*\[\s*([0-9A-Za-z][0-9A-Za-z_.\-]*(?:\s*,\s*[0-9A-Za-z][0-9A-Za-z_.\-]*)*)\s*\]"
+    )
+
+    matches = extras_pattern.match(requirement)
+    if matches:
+        return (
+            matches.group(1),
+            {extra.strip() for extra in matches.group(2).split(",")},
+        )
+
+    return None, None

extract_wheels/lib/requirements_test.py

@@ -0,0 +1,32 @@
+import unittest
+
+from extract_wheels.lib import requirements
+
+
+class TestRequirementExtrasParsing(unittest.TestCase):
+    def test_parses_requirement_for_extra(self) -> None:
+        cases = [
+            ("name[foo]", ("name", frozenset(["foo"]))),
+            ("name[ Foo123 ]", ("name", frozenset(["Foo123"]))),
+            (" name1[ foo ] ", ("name1", frozenset(["foo"]))),
+            (
+                "name [fred,bar] @ http://foo.com ; python_version=='2.7'",
+                ("name", frozenset(["fred", "bar"])),
+            ),
+            (
+                "name[quux, strange];python_version<'2.7' and platform_version=='2'",
+                ("name", frozenset(["quux", "strange"])),
+            ),
+            ("name; (os_name=='a' or os_name=='b') and os_name=='c'", (None, None),),
+            ("name@http://foo.com", (None, None),),
+        ]
+
+        for case, expected in cases:
+            with self.subTest():
+                self.assertTupleEqual(
+                    requirements._parse_requirement_for_extra(case), expected
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

extract_wheels/lib/wheel.py

@@ -2,7 +2,7 @@
 import glob
 import os
 import zipfile
-from typing import Dict, Optional, List, Set
+from typing import Dict, Optional, Set
 
 import pkg_resources
 import pkginfo
@@ -26,7 +26,7 @@ def name(self) -> str:
     def metadata(self) -> pkginfo.Wheel:
         return pkginfo.get_metadata(self.path)
 
-    def dependencies(self, extras_requested: Optional[List[str]] = None) -> Set[str]:
+    def dependencies(self, extras_requested: Optional[Set[str]] = None) -> Set[str]:
         dependency_set = set()
 
         for wheel_req in self.metadata.requires_dist: