Added support for annotating rendered pip dependencies

.bazelrc

@@ -3,8 +3,8 @@
 # This lets us glob() up all the files inside the examples to make them inputs to tests
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, run tools/bazel_integration_test/update_deleted_packages.sh
-build --deleted_packages=examples/build_file_generation,examples/pip_install,examples/pip_parse,examples/py_import,examples/relative_requirements
-query --deleted_packages=examples/build_file_generation,examples/pip_install,examples/pip_parse,examples/py_import,examples/relative_requirements
+build --deleted_packages=examples/build_file_generation,examples/pip_install,examples/pip_parse,examples/pip_repository_annotations,examples/py_import,examples/relative_requirements
+query --deleted_packages=examples/build_file_generation,examples/pip_install,examples/pip_parse,examples/pip_repository_annotations,examples/py_import,examples/relative_requirements
 
 test --test_output=errors
 

BUILD

@@ -32,6 +32,8 @@ filegroup(
         "//python:distribution",
         "//python/pip_install:distribution",
         "//third_party/github.com/bazelbuild/bazel-skylib/lib:distribution",
+        "//third_party/github.com/bazelbuild/bazel-skylib/rules:distribution",
+        "//third_party/github.com/bazelbuild/bazel-skylib/rules/private:distribution",
         "//tools:distribution",
     ],
     visibility = ["//examples:__pkg__"],

docs/pip.md

@@ -35,6 +35,33 @@ It also generates two targets for running pip-compile:
 | kwargs |  other bazel attributes passed to the "_test" rule   |  none |
 
 
+<a name="#package_annotation"></a>
+
+## package_annotation
+
+<pre>
+package_annotation(<a href="#package_annotation-additive_build_content">additive_build_content</a>, <a href="#package_annotation-copy_files">copy_files</a>, <a href="#package_annotation-copy_executables">copy_executables</a>, <a href="#package_annotation-data">data</a>, <a href="#package_annotation-data_exclude_glob">data_exclude_glob</a>,
+                   <a href="#package_annotation-srcs_exclude_glob">srcs_exclude_glob</a>)
+</pre>
+
+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
+
+
+**PARAMETERS**
+
+
+| Name  | Description | Default Value |
+| :-------------: | :-------------: | :-------------: |
+| additive_build_content |  Raw text to add to the generated <code>BUILD</code> file of a package.   |  <code>None</code> |
+| copy_files |  A mapping of <code>src</code> and <code>out</code> files for [@bazel_skylib//rules:copy_file.bzl][cf]   |  <code>{}</code> |
+| copy_executables |  A mapping of <code>src</code> and <code>out</code> files for     [@bazel_skylib//rules:copy_file.bzl][cf]. Targets generated here will also be flagged as     executable.   |  <code>{}</code> |
+| data |  A list of labels to add as <code>data</code> dependencies to the generated <code>py_library</code> target.   |  <code>[]</code> |
+| data_exclude_glob |  A list of exclude glob patterns to add as <code>data</code> to the generated     <code>py_library</code> target.   |  <code>[]</code> |
+| srcs_exclude_glob |  A list of labels to add as <code>srcs</code> to the generated <code>py_library</code> target.   |  <code>[]</code> |
+
+
 <a name="#pip_install"></a>
 
 ## pip_install

examples/BUILD

@@ -27,6 +27,11 @@ bazel_integration_test(
     timeout = "long",
 )
 
+bazel_integration_test(
+    name = "pip_repository_annotations_example",
+    timeout = "long",
+)
+
 bazel_integration_test(
     name = "py_import_example",
     timeout = "long",

examples/pip_repository_annotations/.bazelrc

@@ -0,0 +1,2 @@
+# https://docs.bazel.build/versions/main/best-practices.html#using-the-bazelrc-file
+try-import %workspace%/user.bazelrc

examples/pip_repository_annotations/BUILD

@@ -0,0 +1,30 @@
+load("@pip_installed//:requirements.bzl", "requirement")
+load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:pip.bzl", "compile_pip_requirements")
+
+exports_files(
+    glob(["data/**"]),
+    visibility = ["//visibility:public"],
+)
+
+# This rule adds a convenient way to update the requirements file.
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = ["--allow-unsafe"],
+)
+
+py_test(
+    name = "pip_parse_annotations_test",
+    srcs = ["pip_repository_annotations_test.py"],
+    env = {"WHEEL_PKG_DIR": "pip_parsed_wheel"},
+    main = "pip_repository_annotations_test.py",
+    deps = ["@pip_parsed_wheel//:pkg"],
+)
+
+py_test(
+    name = "pip_install_annotations_test",
+    srcs = ["pip_repository_annotations_test.py"],
+    env = {"WHEEL_PKG_DIR": "pip_installed/pypi__wheel"},
+    main = "pip_repository_annotations_test.py",
+    deps = [requirement("wheel")],
+)

examples/pip_repository_annotations/WORKSPACE

@@ -0,0 +1,58 @@
+workspace(name = "pip_repository_annotations_example")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "rules_python",
+    sha256 = "cd6730ed53a002c56ce4e2f396ba3b3be262fd7cb68339f0377a45e8227fe332",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.5.0/rules_python-0.5.0.tar.gz",
+)
+
+http_archive(
+    name = "bazel_skylib",
+    sha256 = "c6966ec828da198c5d9adbaa94c05e3a1c7f21bd012a0b29ba8ddbccb2c93b0d",
+    urls = [
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.1.1/bazel-skylib-1.1.1.tar.gz",
+        "https://mirror.bazel.build/github.com/bazelbuild/bazel-skylib/releases/download/1.1.1/bazel-skylib-1.1.1.tar.gz",
+    ],
+)
+
+load("@rules_python//python:pip.bzl", "package_annotation", "pip_install", "pip_parse")
+
+# Here we can see an example of annotations being applied to an arbitrary
+# package. For details on `package_annotation` and it's uses, see the
+# docs at @rules_python//docs:pip.md`.
+ANNOTATIONS = {
+    "wheel": package_annotation(
+        additive_build_content = """\
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
+write_file(
+    name = "generated_file",
+    out = "generated_file.txt",
+    content = ["Hello world from build content file"],
+)
+""",
+        copy_executables = {"@pip_repository_annotations_example//:data/copy_executable.py": "copied_content/executable.py"},
+        copy_files = {"@pip_repository_annotations_example//:data/copy_file.txt": "copied_content/file.txt"},
+        data = [":generated_file"],
+        data_exclude_glob = ["*.dist-info/RECORD"],
+    ),
+}
+
+# For a more thorough example of `pip_parse`. See `@rules_python//examples/pip_parse`
+pip_parse(
+    name = "pip_parsed",
+    annotations = ANNOTATIONS,
+    requirements_lock = "//:requirements.txt",
+)
+
+load("@pip_parsed//:requirements.bzl", "install_deps")
+
+install_deps()
+
+# For a more thorough example of `pip_install`. See `@rules_python//examples/pip_install`
+pip_install(
+    name = "pip_installed",
+    annotations = ANNOTATIONS,
+    requirements = "//:requirements.txt",
+)

examples/pip_repository_annotations/data/copy_executable.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python
+
+if __name__ == "__main__":
+    print("Hello world from copied executable")

examples/pip_repository_annotations/data/copy_file.txt

@@ -0,0 +1 @@
+Hello world from copied file

examples/pip_repository_annotations/pip_repository_annotations_test.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+
+import os
+import subprocess
+import unittest
+from glob import glob
+from pathlib import Path
+
+
+class PipRepositoryAnnotationsTest(unittest.TestCase):
+    maxDiff = None
+
+    def wheel_pkg_dir(self) -> str:
+        env = os.environ.get("WHEEL_PKG_DIR")
+        self.assertIsNotNone(env)
+        return env
+
+    def test_build_content_and_data(self):
+        generated_file = (
+            Path.cwd() / "external" / self.wheel_pkg_dir() / "generated_file.txt"
+        )
+        self.assertTrue(generated_file.exists())
+
+        content = generated_file.read_text().rstrip()
+        self.assertEqual(content, "Hello world from build content file")
+
+    def test_copy_files(self):
+        copied_file = (
+            Path.cwd() / "external" / self.wheel_pkg_dir() / "copied_content/file.txt"
+        )
+        self.assertTrue(copied_file.exists())
+
+        content = copied_file.read_text().rstrip()
+        self.assertEqual(content, "Hello world from copied file")
+
+    def test_copy_executables(self):
+        executable = (
+            Path.cwd()
+            / "external"
+            / self.wheel_pkg_dir()
+            / "copied_content/executable.py"
+        )
+        self.assertTrue(executable.exists())
+
+        proc = subprocess.run(
+            [executable], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+        stdout = proc.stdout.decode("utf-8").strip()
+        self.assertEqual(stdout, "Hello world from copied executable")
+
+    def test_data_exclude_glob(self):
+        files = glob("external/" + self.wheel_pkg_dir() + "/wheel-*.dist-info/*")
+        basenames = [Path(path).name for path in files]
+        self.assertIn("WHEEL", basenames)
+        self.assertNotIn("RECORD", basenames)
+
+
+if __name__ == "__main__":
+    unittest.main()

examples/pip_repository_annotations/requirements.in

@@ -0,0 +1 @@
+wheel

examples/pip_repository_annotations/requirements.txt

@@ -0,0 +1,10 @@
+#
+# This file is autogenerated by pip-compile
+# To update, run:
+#
+#    bazel run //:requirements.update
+#
+wheel==0.37.1 \
+    --hash=sha256:4bdcd7d840138086126cd09254dc6195fb4fc6f01c050a1d7236f2630db1d22a \
+    --hash=sha256:e9a504e793efbca1b8e0e9cb979a249cf4a0a7b5b8c9e8b65a5e39d49529c1c4
+    # via -r requirements.in

python/pip.bzl

@@ -13,11 +13,12 @@
 # limitations under the License.
 """Import pip requirements into Bazel."""
 
-load("//python/pip_install:pip_repository.bzl", "pip_repository")
+load("//python/pip_install:pip_repository.bzl", "pip_repository", _package_annotation = "package_annotation")
 load("//python/pip_install:repositories.bzl", "pip_install_dependencies")
 load("//python/pip_install:requirements.bzl", _compile_pip_requirements = "compile_pip_requirements")
 
 compile_pip_requirements = _compile_pip_requirements
+package_annotation = _package_annotation
 
 def pip_install(requirements, name = "pip", **kwargs):
     """Accepts a `requirements.txt` file and installs the dependencies listed within.

python/pip_install/extract_wheels/__init__.py

@@ -12,7 +12,13 @@
 import subprocess
 import sys
 
-from python.pip_install.extract_wheels.lib import arguments, bazel, requirements
+from python.pip_install.extract_wheels.lib import (
+    annotation,
+    arguments,
+    bazel,
+    requirements,
+    wheel,
+)
 
 
 def configure_reproducible_wheels() -> None:
@@ -58,6 +64,11 @@ def main() -> None:
         required=True,
         help="Path to requirements.txt from where to install dependencies",
     )
+    parser.add_argument(
+        "--annotations",
+        type=annotation.annotations_map_from_str_path,
+        help="A json encoded file containing annotations for rendered packages.",
+    )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
     deserialized_args = dict(vars(args))
@@ -89,18 +100,26 @@ def main() -> None:
 
     repo_label = "@%s" % args.repo
 
+    # Locate all wheels
+    wheels = [whl for whl in glob.glob("*.whl")]
+
+    # Collect all annotations
+    reqs = {whl: wheel.Wheel(whl).name for whl in wheels}
+    annotations = args.annotations.collect(reqs.values())
+
     targets = [
         '"{}{}"'.format(
             repo_label,
             bazel.extract_wheel(
-                whl,
-                extras,
-                deserialized_args["pip_data_exclude"],
-                args.enable_implicit_namespace_pkgs,
-                args.repo_prefix,
+                wheel_file=whl,
+                extras=extras,
+                pip_data_exclude=deserialized_args["pip_data_exclude"],
+                enable_implicit_namespace_pkgs=args.enable_implicit_namespace_pkgs,
+                repo_prefix=args.repo_prefix,
+                annotation=annotations.get(name),
             ),
         )
-        for whl in glob.glob("*.whl")
+        for whl, name in reqs.items()
     ]
 
     with open("requirements.bzl", "w") as requirement_file:

python/pip_install/extract_wheels/lib/BUILD

@@ -1,9 +1,11 @@
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
 load("//python/pip_install:repositories.bzl", "requirement")
+load(":annotations_test_helpers.bzl", "package_annotation", "package_annotations_file")
 
 py_library(
     name = "lib",
     srcs = [
+        "annotation.py",
         "arguments.py",
         "bazel.py",
         "namespace_pkgs.py",
@@ -21,6 +23,43 @@ py_library(
     ],
 )
 
+package_annotations_file(
+    name = "mock_annotations",
+    annotations = {
+        "pkg_a": package_annotation(),
+        "pkg_b": package_annotation(
+            data_exclude_glob = [
+                "*.foo",
+                "*.bar",
+            ],
+        ),
+        "pkg_c": package_annotation(
+            additive_build_content = """\
+cc_library(
+    name = "my_target",
+    hdrs = glob(["**/*.h"]),
+    srcs = glob(["**/*.cc"]),
+)
+""",
+            data = [":my_target"],
+        ),
+        "pkg_d": package_annotation(
+            srcs_exclude_glob = ["pkg_d/tests/**"],
+        ),
+    },
+    tags = ["manual"],
+)
+
+py_test(
+    name = "annotations_test",
+    size = "small",
+    srcs = ["annotations_test.py"],
+    data = [":mock_annotations"],
+    env = {"MOCK_ANNOTATIONS": "$(rootpath :mock_annotations)"},
+    tags = ["unit"],
+    deps = [":lib"],
+)
+
 py_test(
     name = "bazel_test",
     size = "small",