Allow users to inject extra deps into pip packages

python/pip_install/extract_wheels/BUILD

@@ -61,6 +61,9 @@ cc_library(
         "pkg_d": package_annotation(
             srcs_exclude_glob = ["pkg_d/tests/**"],
         ),
+        "pkg_e": package_annotation(
+            deps = ["//some:dep"],
+        ),
     },
     tags = ["manual"],
 )

python/pip_install/extract_wheels/annotation.py

@@ -19,6 +19,7 @@ def __init__(self, content: Dict[str, Any]) -> None:
             "data",
             "data_exclude_glob",
             "srcs_exclude_glob",
+            "deps",
         ):
             if field not in content:
                 missing.append(field)
@@ -61,6 +62,10 @@ def data_exclude_glob(self) -> List[str]:
     def srcs_exclude_glob(self) -> List[str]:
         return self["srcs_exclude_glob"]
 
+    @property
+    def deps(self) -> List[str]:
+        return self["deps"]
+
 
 class AnnotationsMap:
     """A mapping of python package names to [Annotation]"""

python/pip_install/extract_wheels/annotations_test.py

@@ -25,10 +25,12 @@ def test_annotations_constructor(self) -> None:
         annotations_map = AnnotationsMap(annotations_path)
         self.assertListEqual(
             list(annotations_map.annotations.keys()),
-            ["pkg_a", "pkg_b", "pkg_c", "pkg_d"],
+            ["pkg_a", "pkg_b", "pkg_c", "pkg_d", "pkg_e"],
         )
 
-        collection = annotations_map.collect(["pkg_a", "pkg_b", "pkg_c", "pkg_d"])
+        collection = annotations_map.collect(
+            ["pkg_a", "pkg_b", "pkg_c", "pkg_d", "pkg_e"]
+        )
 
         self.assertEqual(
             collection["pkg_a"],
@@ -40,6 +42,7 @@ def test_annotations_constructor(self) -> None:
                     "data": [],
                     "data_exclude_glob": [],
                     "srcs_exclude_glob": [],
+                    "deps": [],
                 }
             ),
         )
@@ -54,6 +57,7 @@ def test_annotations_constructor(self) -> None:
                     "data": [],
                     "data_exclude_glob": ["*.foo", "*.bar"],
                     "srcs_exclude_glob": [],
+                    "deps": [],
                 }
             ),
         )
@@ -84,6 +88,7 @@ def test_annotations_constructor(self) -> None:
                     "data": [":my_target"],
                     "data_exclude_glob": [],
                     "srcs_exclude_glob": [],
+                    "deps": [],
                 }
             ),
         )
@@ -98,6 +103,22 @@ def test_annotations_constructor(self) -> None:
                     "data": [],
                     "data_exclude_glob": [],
                     "srcs_exclude_glob": ["pkg_d/tests/**"],
+                    "deps": [],
+                }
+            ),
+        )
+
+        self.assertEqual(
+            collection["pkg_e"],
+            Annotation(
+                {
+                    "additive_build_content": None,
+                    "copy_executables": {},
+                    "copy_files": {},
+                    "data": [],
+                    "data_exclude_glob": [],
+                    "srcs_exclude_glob": [],
+                    "deps": ["//some:dep"],
                 }
             ),
         )

python/pip_install/extract_wheels/bazel.py

@@ -415,6 +415,7 @@ def extract_wheel(
         data = []
         data_exclude = pip_data_exclude
         srcs_exclude = []
+        deps = sanitised_dependencies
         if annotation:
             for src, dest in annotation.copy_files.items():
                 data.append(dest)
@@ -427,14 +428,15 @@ def extract_wheel(
             data.extend(annotation.data)
             data_exclude.extend(annotation.data_exclude_glob)
             srcs_exclude.extend(annotation.srcs_exclude_glob)
+            deps.extend('"%s"' % dep for dep in annotation.deps)
             if annotation.additive_build_content:
                 additional_content.append(annotation.additive_build_content)
 
         contents = generate_build_file_contents(
             name=PY_LIBRARY_LABEL
             if incremental
             else sanitise_name(whl.name, repo_prefix),
-            dependencies=sanitised_dependencies,
+            dependencies=deps,
             whl_file_deps=sanitised_wheel_file_dependencies,
             data_exclude=data_exclude,
             data=data,

python/pip_install/extract_wheels/bazel_test.py

@@ -1,11 +1,80 @@
+import os
+import shutil
 import unittest
+import unittest.mock
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
 
-from python.pip_install.extract_wheels.bazel import generate_entry_point_contents
+from python.pip_install.extract_wheels import annotation
+from python.pip_install.extract_wheels.bazel import (
+    extract_wheel,
+    generate_entry_point_contents,
+)
+
+
+class MockWheelInstance:
+    """A mock for python.pip_install.extract_wheels.wheel.Wheel."""
+
+    def __init__(
+        self,
+        name: str,
+        version: str,
+        path: str,
+        dependencies: List[str],
+        entry_points: Dict[str, Tuple[str, str]],
+    ):
+        self.name = name
+        self.path = path
+        self.version = version
+        self.dependency_list = dependencies
+        self.entry_point_dict = entry_points
+
+    def dependencies(self, extras_requested):
+        # Since the caller can customize the dependency list to their liking,
+        # we don't need to act on the extras_requested.
+        return set(self.dependency_list)
+
+    def entry_points(self):
+        return self.entry_point_dict
+
+    def unzip(self, directory):
+        # We don't care about actually generating any files for our purposes.
+        pass
+
+
+def parse_starlark(filename: os.PathLike) -> Dict[str, unittest.mock.MagicMock]:
+    """Parses a Starlark file.
+
+    Args:
+        filename: The name of the file to parse as Starlark.
+
+    Returns:
+        A dictionary of MagicMock instances for each of the functions that was
+        invoked in the Starlark file.
+    """
+    starlark_globals = {
+        "filegroup": unittest.mock.MagicMock(),
+        "glob": unittest.mock.MagicMock(return_value=["<glob()>"]),
+        "load": unittest.mock.MagicMock(),
+        "package": unittest.mock.MagicMock(),
+        "py_binary": unittest.mock.MagicMock(),
+        "py_library": unittest.mock.MagicMock(),
+    }
+    compiled_starlark = compile(Path(filename).read_text(), filename, "exec")
+    eval(compiled_starlark, starlark_globals)
+    return starlark_globals
 
 
 class BazelTestCase(unittest.TestCase):
+    def setUp(self):
+        self.tmpdir = Path(os.environ["TEST_TMPDIR"]) / "tmpdir"
+        self.tmpdir.mkdir()
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir)
+
     def test_generate_entry_point_contents(self):
-        got = generate_entry_point_contents("sphinx.cmd.build:main")
+        got = generate_entry_point_contents("sphinx.cmd.build", "main")
         want = """#!/usr/bin/env python3
 import sys
 from sphinx.cmd.build import main
@@ -16,11 +85,111 @@ def test_generate_entry_point_contents(self):
 
     def test_generate_entry_point_contents_with_shebang(self):
         got = generate_entry_point_contents(
-            "sphinx.cmd.build:main", shebang="#!/usr/bin/python"
+            "sphinx.cmd.build", "main", shebang="#!/usr/bin/python"
         )
         want = """#!/usr/bin/python
 import sys
 from sphinx.cmd.build import main
-sys.exit(main())
+if __name__ == "__main__":
+    sys.exit(main())
 """
         self.assertEqual(got, want)
+
+    @unittest.mock.patch("python.pip_install.extract_wheels.wheel.Wheel")
+    def test_extract_wheel(self, MockWheel):
+        """Validates that extract_wheel generates the expected BUILD file.
+
+        We don't really care about extracting a .whl file here so we mock that
+        part. The interesting bit is the BUILD file generation.
+        """
+        # Create a dummy wheel that we pretend to extract.
+        mock_wheel_instance = MockWheelInstance(
+            name="test-wheel",
+            version="1.2.3",
+            path="path/to/test-wheel.whl",
+            dependencies=["a", "b"],
+            entry_points={
+                "test_bin_entry": ("test_wheel.entry", "main"),
+            },
+        )
+        MockWheel.return_value = mock_wheel_instance
+
+        # Run the BUILD file generation code.
+        extract_wheel(
+            wheel_file=mock_wheel_instance.path,
+            extras={},
+            pip_data_exclude=[],
+            enable_implicit_namespace_pkgs=True,
+            repo_prefix="repo_prefix_",
+            incremental=True,
+            incremental_dir=self.tmpdir,
+            annotation=annotation.Annotation(
+                {
+                    "additive_build_content": [],
+                    "copy_executables": {},
+                    "copy_files": {},
+                    "data": ["//some/extra:data"],
+                    "data_exclude_glob": ["foo/bad.data.*"],
+                    "srcs_exclude_glob": ["foo/bad.srcs.*"],
+                    "deps": ["//a/dep/of:some_kind"],
+                }
+            ),
+        )
+
+        parsed_starlark = parse_starlark(self.tmpdir / "BUILD.bazel")
+
+        # Validate the library target.
+        self.assertListEqual(
+            parsed_starlark["py_library"].mock_calls,
+            [
+                unittest.mock.call(
+                    name="pkg",
+                    srcs=["<glob()>"],
+                    data=["//some/extra:data", "<glob()>"],
+                    imports=["site-packages"],
+                    deps=[
+                        "//a/dep/of:some_kind",
+                        "@repo_prefix_a//:pkg",
+                        "@repo_prefix_b//:pkg",
+                    ],
+                    tags=["pypi_name=test-wheel", "pypi_version=1.2.3"],
+                ),
+            ],
+        )
+        self.assertListEqual(
+            parsed_starlark["glob"].mock_calls[3:],
+            [
+                unittest.mock.call(
+                    ["site-packages/**/*.py"],
+                    exclude=["foo/bad.srcs.*"],
+                    allow_empty=True,
+                ),
+                unittest.mock.call(
+                    ["site-packages/**/*"],
+                    exclude=[
+                        "**/* *",
+                        "**/*.dist-info/RECORD",
+                        "**/*.py",
+                        "**/*.pyc",
+                        "foo/bad.data.*",
+                    ],
+                ),
+            ],
+        )
+
+        # Validate the entry point targets.
+        self.assertListEqual(
+            parsed_starlark["py_binary"].mock_calls,
+            [
+                unittest.mock.call(
+                    name="rules_python_wheel_entry_point_test_bin_entry",
+                    srcs=["rules_python_wheel_entry_point_test_bin_entry.py"],
+                    imports=["."],
+                    deps=["pkg"],
+                ),
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

python/pip_install/pip_repository.bzl

@@ -508,7 +508,8 @@ def package_annotation(
         copy_executables = {},
         data = [],
         data_exclude_glob = [],
-        srcs_exclude_glob = []):
+        srcs_exclude_glob = [],
+        deps = []):
     """Annotations to apply to the BUILD file content from package generated from a `pip_repository` rule.
 
     [cf]: https://github.com/bazelbuild/bazel-skylib/blob/main/docs/copy_file_doc.md
@@ -523,6 +524,7 @@ def package_annotation(
         data_exclude_glob (list, optional): A list of exclude glob patterns to add as `data` to the generated
             `py_library` target.
         srcs_exclude_glob (list, optional): A list of labels to add as `srcs` to the generated `py_library` target.
+        deps (list, optional): A list of labels to add as `deps` to the generated `py_library` target.
 
     Returns:
         str: A json encoded string of the provided content.
@@ -534,4 +536,5 @@ def package_annotation(
         data = data,
         data_exclude_glob = data_exclude_glob,
         srcs_exclude_glob = srcs_exclude_glob,
+        deps = deps,
     ))