feat(bzlmod): Implementing wheel annotations via whl_mods (bazel-contrib#1278)

chrislovecnm · web-flow · commit 42b72dbd6ea3 · 2023-07-10T15:55:46.000Z
This commit implements a bzlmod extension that allows users to create "annotations" for wheel builds. The wheel_builder.py accepts a JSON file via a parameter called annotations; this extension creates those JSON files. The pip extension accepts a Label -> String dict argument of the JSON files. This feature is renamed to `whl_mods` because the JSON files are handled differently and the name "annotations" is uninformative. This modifies the creation of the BUILD files and their content, and is much more than just adding some notes about a whl. The whl_mod extension wheel names and the wheel names in pip must match. Closes: bazel-contrib#1213
diff --git a/.bazelrc b/.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/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_point,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_point,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
 
 test --test_output=errors
 
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
@@ -37,8 +37,50 @@ python.toolchain(
 # rules based on the `python_version` arg values.
 use_repo(python, "python_3_10", "python_3_9", "python_aliases")
 
+# This extension allows a user to create modifications to how rules_python
+# creates different wheel repositories.  Different attributes allow the user
+# to modify the BUILD file, and copy files.
+# See @rules_python//python/extensions:whl_mods.bzl attributes for more information
+# on each of the attributes.
+# You are able to set a hub name, so that you can have different modifications of the same
+# wheel in different pip hubs.
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 
+# Call whl_mods.create for the requests package.
+pip.whl_mods(
+    # we are using the appended_build_content.BUILD file
+    # to add content to the request wheel BUILD file.
+    additive_build_content_file = "//whl_mods:appended_build_content.BUILD",
+    data = [":generated_file"],
+    hub_name = "whl_mods_hub",
+    whl_name = "requests",
+)
+
+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"],
+)
+"""
+
+# Call whl_mods.create for the wheel package.
+pip.whl_mods(
+    additive_build_content = ADDITIVE_BUILD_CONTENT,
+    copy_executables = {
+        "@@//whl_mods:data/copy_executable.py": "copied_content/executable.py",
+    },
+    copy_files = {
+        "@@//whl_mods:data/copy_file.txt": "copied_content/file.txt",
+    },
+    data = [":generated_file"],
+    data_exclude_glob = ["site-packages/*.dist-info/WHEEL"],
+    hub_name = "whl_mods_hub",
+    whl_name = "wheel",
+)
+use_repo(pip, "whl_mods_hub")
+
 # To fetch pip dependencies, use pip.parse. We can pass in various options,
 # but typically we pass requirements and the Python version. The Python
 # version must have been configured by a corresponding `python.toolchain()`
@@ -50,12 +92,26 @@ pip.parse(
     python_version = "3.9",
     requirements_lock = "//:requirements_lock_3_9.txt",
     requirements_windows = "//:requirements_windows_3_9.txt",
+    # These modifications were created above and we
+    # are providing pip.parse with the label of the mod
+    # and the name of the wheel.
+    whl_modifications = {
+        "@whl_mods_hub//:requests.json": "requests",
+        "@whl_mods_hub//:wheel.json": "wheel",
+    },
 )
 pip.parse(
     hub_name = "pip",
     python_version = "3.10",
     requirements_lock = "//:requirements_lock_3_10.txt",
     requirements_windows = "//:requirements_windows_3_10.txt",
+    # These modifications were created above and we
+    # are providing pip.parse with the label of the mod
+    # and the name of the wheel.
+    whl_modifications = {
+        "@whl_mods_hub//:requests.json": "requests",
+        "@whl_mods_hub//:wheel.json": "wheel",
+    },
 )
 
 # NOTE: The pip_39 repo is only used because the plain `@pip` repo doesn't
diff --git a/examples/bzlmod/requirements.in b/examples/bzlmod/requirements.in
@@ -1,3 +1,6 @@
+--extra-index-url https://pypi.python.org/simple/
+
+wheel
 websockets
 requests~=2.25.1
 s3cmd~=2.1.0
diff --git a/examples/bzlmod/requirements_lock_3_10.txt b/examples/bzlmod/requirements_lock_3_10.txt
@@ -4,6 +4,8 @@
 #
 #    bazel run //:requirements_3_10.update
 #
+--extra-index-url https://pypi.python.org/simple/
+
 astroid==2.13.5 \
     --hash=sha256:6891f444625b6edb2ac798829b689e95297e100ddf89dbed5a8c610e34901501 \
     --hash=sha256:df164d5ac811b9f44105a72b8f9d5edfb7b5b2d7e979b04ea377a77b3229114a
@@ -238,6 +240,10 @@ websockets==11.0.3 \
     --hash=sha256:fb06eea71a00a7af0ae6aefbb932fb8a7df3cb390cc217d51a9ad7343de1b8d0 \
     --hash=sha256:ffd7dcaf744f25f82190856bc26ed81721508fc5cbf2a330751e135ff1283564
     # via -r requirements.in
+wheel==0.40.0 \
+    --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \
+    --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247
+    # via -r requirements.in
 wrapt==1.15.0 \
     --hash=sha256:02fce1852f755f44f95af51f69d22e45080102e9d00258053b79367d07af39c0 \
     --hash=sha256:077ff0d1f9d9e4ce6476c1a924a3332452c1406e59d90a2cf24aeb29eeac9420 \
diff --git a/examples/bzlmod/requirements_lock_3_9.txt b/examples/bzlmod/requirements_lock_3_9.txt
@@ -4,6 +4,8 @@
 #
 #    bazel run //:requirements_3_9.update
 #
+--extra-index-url https://pypi.python.org/simple/
+
 astroid==2.12.13 \
     --hash=sha256:10e0ad5f7b79c435179d0d0f0df69998c4eef4597534aae44910db060baeb907 \
     --hash=sha256:1493fe8bd3dfd73dc35bd53c9d5b6e49ead98497c47b2307662556a5692d29d7
@@ -227,6 +229,10 @@ websockets==11.0.3 \
     --hash=sha256:fb06eea71a00a7af0ae6aefbb932fb8a7df3cb390cc217d51a9ad7343de1b8d0 \
     --hash=sha256:ffd7dcaf744f25f82190856bc26ed81721508fc5cbf2a330751e135ff1283564
     # via -r requirements.in
+wheel==0.40.0 \
+    --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \
+    --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247
+    # via -r requirements.in
 wrapt==1.14.1 \
     --hash=sha256:00b6d4ea20a906c0ca56d84f93065b398ab74b927a7a3dbd470f6fc503f95dc3 \
     --hash=sha256:01c205616a89d09827986bc4e859bcabd64f5a0662a7fe95e0d359424e0e071b \
diff --git a/examples/bzlmod/requirements_windows_3_10.txt b/examples/bzlmod/requirements_windows_3_10.txt
@@ -4,6 +4,8 @@
 #
 #    bazel run //:requirements_3_10.update
 #
+--extra-index-url https://pypi.python.org/simple/
+
 astroid==2.13.5 \
     --hash=sha256:6891f444625b6edb2ac798829b689e95297e100ddf89dbed5a8c610e34901501 \
     --hash=sha256:df164d5ac811b9f44105a72b8f9d5edfb7b5b2d7e979b04ea377a77b3229114a
@@ -242,6 +244,10 @@ websockets==11.0.3 \
     --hash=sha256:fb06eea71a00a7af0ae6aefbb932fb8a7df3cb390cc217d51a9ad7343de1b8d0 \
     --hash=sha256:ffd7dcaf744f25f82190856bc26ed81721508fc5cbf2a330751e135ff1283564
     # via -r requirements.in
+wheel==0.40.0 \
+    --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \
+    --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247
+    # via -r requirements.in
 wrapt==1.15.0 \
     --hash=sha256:02fce1852f755f44f95af51f69d22e45080102e9d00258053b79367d07af39c0 \
     --hash=sha256:077ff0d1f9d9e4ce6476c1a924a3332452c1406e59d90a2cf24aeb29eeac9420 \
diff --git a/examples/bzlmod/requirements_windows_3_9.txt b/examples/bzlmod/requirements_windows_3_9.txt
@@ -4,6 +4,8 @@
 #
 #    bazel run //:requirements_3_9.update
 #
+--extra-index-url https://pypi.python.org/simple/
+
 astroid==2.12.13 \
     --hash=sha256:10e0ad5f7b79c435179d0d0f0df69998c4eef4597534aae44910db060baeb907 \
     --hash=sha256:1493fe8bd3dfd73dc35bd53c9d5b6e49ead98497c47b2307662556a5692d29d7
@@ -231,6 +233,10 @@ websockets==11.0.3 \
     --hash=sha256:fb06eea71a00a7af0ae6aefbb932fb8a7df3cb390cc217d51a9ad7343de1b8d0 \
     --hash=sha256:ffd7dcaf744f25f82190856bc26ed81721508fc5cbf2a330751e135ff1283564
     # via -r requirements.in
+wheel==0.40.0 \
+    --hash=sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873 \
+    --hash=sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247
+    # via -r requirements.in
 wrapt==1.14.1 \
     --hash=sha256:00b6d4ea20a906c0ca56d84f93065b398ab74b927a7a3dbd470f6fc503f95dc3 \
     --hash=sha256:01c205616a89d09827986bc4e859bcabd64f5a0662a7fe95e0d359424e0e071b \
diff --git a/examples/bzlmod/whl_mods/BUILD.bazel b/examples/bzlmod/whl_mods/BUILD.bazel
@@ -0,0 +1,21 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+exports_files(
+    glob(["data/**"]),
+    visibility = ["//visibility:public"],
+)
+
+py_test(
+    name = "pip_whl_mods_test",
+    srcs = ["pip_whl_mods_test.py"],
+    env = {
+        "REQUESTS_PKG_DIR": "pip_39_requests",
+        "WHEEL_PKG_DIR": "pip_39_wheel",
+    },
+    main = "pip_whl_mods_test.py",
+    deps = [
+        "@pip//requests:pkg",
+        "@pip//wheel:pkg",
+        "@rules_python//python/runfiles",
+    ],
+)
diff --git a/examples/bzlmod/whl_mods/appended_build_content.BUILD b/examples/bzlmod/whl_mods/appended_build_content.BUILD
@@ -0,0 +1,7 @@
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
+
+write_file(
+    name = "generated_file",
+    out = "generated_file.txt",
+    content = ["Hello world from requests"],
+)
diff --git a/examples/bzlmod/whl_mods/data/copy_executable.py b/examples/bzlmod/whl_mods/data/copy_executable.py
@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+if __name__ == "__main__":
+    print("Hello world from copied executable")
diff --git a/examples/bzlmod/whl_mods/data/copy_file.txt b/examples/bzlmod/whl_mods/data/copy_file.txt
@@ -0,0 +1 @@
+Hello world from copied file
diff --git a/examples/bzlmod/whl_mods/pip_whl_mods_test.py b/examples/bzlmod/whl_mods/pip_whl_mods_test.py
@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import platform
+import subprocess
+import sys
+import unittest
+from pathlib import Path
+
+from python.runfiles import runfiles
+
+
+class PipWhlModsTest(unittest.TestCase):
+    maxDiff = None
+
+    def package_path(self) -> str:
+        return "rules_python~override~pip~"
+
+    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):
+        r = runfiles.Create()
+        rpath = r.Rlocation(
+                "{}{}/generated_file.txt".format(
+                    self.package_path(), 
+                    self.wheel_pkg_dir(),
+                    ),
+                )
+        generated_file = Path(rpath)
+        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):
+        r = runfiles.Create()
+        rpath = r.Rlocation(
+                "{}{}/copied_content/file.txt".format(
+                    self.package_path(), 
+                    self.wheel_pkg_dir(),
+                    )
+                )
+        copied_file = Path(rpath)
+        self.assertTrue(copied_file.exists())
+
+        content = copied_file.read_text().rstrip()
+        self.assertEqual(content, "Hello world from copied file")
+
+    def test_copy_executables(self):
+        r = runfiles.Create()
+        rpath = r.Rlocation(
+                "{}{}/copied_content/executable{}".format(
+                    self.package_path(), 
+                    self.wheel_pkg_dir(),
+                    ".exe" if platform.system() == "windows" else ".py",
+                    )
+                )
+        executable = Path(rpath)
+        self.assertTrue(executable.exists())
+
+        proc = subprocess.run(
+            [sys.executable, str(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):
+        current_wheel_version = "0.40.0"
+
+        r = runfiles.Create()
+        dist_info_dir = "{}{}/site-packages/wheel-{}.dist-info".format(
+                self.package_path(), 
+                self.wheel_pkg_dir(),
+                current_wheel_version,
+                )
+
+        # Note: `METADATA` is important as it's consumed by https://docs.python.org/3/library/importlib.metadata.html
+        # `METADATA` is expected to be there to show dist-info files are included in the runfiles.
+        metadata_path = r.Rlocation("{}/METADATA".format(dist_info_dir))
+
+        # However, `WHEEL` was explicitly excluded, so it should be missing
+        wheel_path = r.Rlocation("{}/WHEEL".format(dist_info_dir))
+
+        self.assertTrue(Path(metadata_path).exists())
+        self.assertFalse(Path(wheel_path).exists())
+
+    def requests_pkg_dir(self) -> str:
+        env = os.environ.get("REQUESTS_PKG_DIR")
+        self.assertIsNotNone(env)
+        return env
+
+    def test_extra(self):
+        # This test verifies that annotations work correctly for pip packages with extras
+        # specified, in this case requests[security].
+        r = runfiles.Create()
+        rpath = r.Rlocation(
+                "{}{}/generated_file.txt".format(
+                    self.package_path(), 
+                    self.requests_pkg_dir(),
+                    ),
+                )
+        generated_file = Path(rpath)
+        self.assertTrue(generated_file.exists())
+
+        content = generated_file.read_text().rstrip()
+        self.assertEqual(content, "Hello world from requests")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/python/extensions/pip.bzl b/python/extensions/pip.bzl
