feat(whl_library): generate platform-specific dependency closures (#1593)

Before this change, the dependency closures would be influenced by the
host-python interpreter, this removes the influence by detecting the
platforms against which the `Requires-Dist` wheel metadata is evaluated.
This functionality can be enabled via `experimental_target_platforms`
attribute to the `pip.parse` extension and is showcased in the `bzlmod`
example. The same attribute is also supported on the legacy `pip_parse`
repository rule.

The detection works in the following way:
- Check if the python wheel is platform specific or cross-platform
  (i.e., ends with `any.whl`), if it is then platform-specific
  dependencies are generated, which will go through a `select`
  statement.
- If it is platform specific, then parse the platform_tag and evaluate
  the `Requires-Dist` markers assuming the target platform rather than
  the host platform.

NOTE: The `whl` `METADATA` is now being parsed using the `packaging`
Python package instead of `pkg_resources` from `setuptools`.

Fixes #1591

---------

Co-authored-by: Richard Levasseur <[email protected]>

Author: aignas

CHANGELOG.md

@@ -36,6 +36,13 @@ A brief description of the categories of changes:
 * (gazelle) The gazelle plugin helper was not working with Python toolchains 3.11
   and above due to a bug in the helper components not being on PYTHONPATH.
 
+* (pip_parse) The repositories created by `whl_library` can now parse the `whl`
+  METADATA and generate dependency closures irrespective of the host platform
+  the generation is executed on. This can be turned on by supplying
+  `experimental_target_platforms = ["all"]` to the `pip_parse` or the `bzlmod`
+  equivalent. This may help in cases where fetching wheels for a different
+  platform using `download_only = True` feature.
+
 [0.XX.0]: https://github.com/bazelbuild/rules_python/releases/tag/0.XX.0
 
 ## [0.27.0] - 2023-11-16

examples/bzlmod/MODULE.bazel

@@ -98,6 +98,13 @@ pip.parse(
             "sphinxcontrib-serializinghtml",
         ],
     },
+    # You can use one of the values below to specify the target platform
+    # to generate the dependency graph for.
+    experimental_target_platforms = [
+        "all",
+        "linux_*",
+        "host",
+    ],
     hub_name = "pip",
     python_version = "3.9",
     requirements_lock = "//:requirements_lock_3_9.txt",
@@ -121,6 +128,13 @@ pip.parse(
             "sphinxcontrib-serializinghtml",
         ],
     },
+    # You can use one of the values below to specify the target platform
+    # to generate the dependency graph for.
+    experimental_target_platforms = [
+        "all",
+        "linux_*",
+        "host",
+    ],
     hub_name = "pip",
     python_version = "3.10",
     requirements_lock = "//:requirements_lock_3_10.txt",

python/pip_install/pip_repository.bzl

@@ -345,6 +345,8 @@ def _pip_repository_impl(rctx):
 
     if rctx.attr.python_interpreter_target:
         config["python_interpreter_target"] = str(rctx.attr.python_interpreter_target)
+    if rctx.attr.experimental_target_platforms:
+        config["experimental_target_platforms"] = rctx.attr.experimental_target_platforms
 
     if rctx.attr.incompatible_generate_aliases:
         macro_tmpl = "@%s//{}:{}" % rctx.attr.name
@@ -472,6 +474,30 @@ Warning:
   If a dependency participates in multiple cycles, all of those cycles must be
   collapsed down to one. For instance `a <-> b` and `a <-> c` cannot be listed
   as two separate cycles.
+""",
+    ),
+    "experimental_target_platforms": attr.string_list(
+        default = [],
+        doc = """\
+A list of platforms that we will generate the conditional dependency graph for
+cross platform wheels by parsing the wheel metadata. This will generate the
+correct dependencies for packages like `sphinx` or `pylint`, which include
+`colorama` when installed and used on Windows platforms.
+
+An empty list means falling back to the legacy behaviour where the host
+platform is the target platform.
+
+WARNING: It may not work as expected in cases where the python interpreter
+implementation that is being used at runtime is different between different platforms.
+This has been tested for CPython only.
+
+Special values: `all` (for generating deps for all platforms), `host` (for
+generating deps for the host platform only). `linux_*` and other `<os>_*` values.
+In the future we plan to set `all` as the default to this attribute.
+
+For specific target platforms use values of the form `<os>_<arch>` where `<os>`
+is one of `linux`, `osx`, `windows` and arch is one of `x86_64`, `x86_32`,
+`aarch64`, `s390x` and `ppc64le`.
 """,
     ),
     "extra_pip_args": attr.string_list(
@@ -713,7 +739,10 @@ def _whl_library_impl(rctx):
         )
 
     result = rctx.execute(
-        args + ["--whl-file", whl_path],
+        args + [
+            "--whl-file",
+            whl_path,
+        ] + ["--platform={}".format(p) for p in rctx.attr.experimental_target_platforms],
         environment = environment,
         quiet = rctx.attr.quiet,
         timeout = rctx.attr.timeout,
@@ -749,6 +778,7 @@ def _whl_library_impl(rctx):
         repo_prefix = rctx.attr.repo_prefix,
         whl_name = whl_path.basename,
         dependencies = metadata["deps"],
+        dependencies_by_platform = metadata["deps_by_platform"],
         group_name = rctx.attr.group_name,
         group_deps = rctx.attr.group_deps,
         data_exclude = rctx.attr.pip_data_exclude,
@@ -815,7 +845,7 @@ whl_library_attrs = {
         doc = "Python requirement string describing the package to make available",
     ),
     "whl_patches": attr.label_keyed_string_dict(
-        doc = """"a label-keyed-string dict that has
+        doc = """a label-keyed-string dict that has
             json.encode(struct([whl_file], patch_strip]) as values. This
             is to maintain flexibility and correct bzlmod extension interface
             until we have a better way to define whl_library and move whl

python/pip_install/private/generate_whl_library_build_bazel.bzl

@@ -25,6 +25,7 @@ load(
     "WHEEL_FILE_PUBLIC_LABEL",
 )
 load("//python/private:normalize_name.bzl", "normalize_name")
+load("//python/private:text_util.bzl", "render")
 
 _COPY_FILE_TEMPLATE = """\
 copy_file(
@@ -101,11 +102,36 @@ alias(
 )
 """
 
+def _render_list_and_select(deps, deps_by_platform, tmpl):
+    deps = render.list([tmpl.format(d) for d in deps])
+
+    if not deps_by_platform:
+        return deps
+
+    deps_by_platform = {
+        p if p.startswith("@") else ":is_" + p: [
+            tmpl.format(d)
+            for d in deps
+        ]
+        for p, deps in deps_by_platform.items()
+    }
+
+    # Add the default, which means that we will be just using the dependencies in
+    # `deps` for platforms that are not handled in a special way by the packages
+    deps_by_platform["//conditions:default"] = []
+    deps_by_platform = render.select(deps_by_platform, value_repr = render.list)
+
+    if deps == "[]":
+        return deps_by_platform
+    else:
+        return "{} + {}".format(deps, deps_by_platform)
+
 def generate_whl_library_build_bazel(
         *,
         repo_prefix,
         whl_name,
         dependencies,
+        dependencies_by_platform,
         data_exclude,
         tags,
         entry_points,
@@ -118,6 +144,7 @@ def generate_whl_library_build_bazel(
         repo_prefix: the repo prefix that should be used for dependency lists.
         whl_name: the whl_name that this is generated for.
         dependencies: a list of PyPI packages that are dependencies to the py_library.
+        dependencies_by_platform: a dict[str, list] of PyPI packages that may vary by platform.
         data_exclude: more patterns to exclude from the data attribute of generated py_library rules.
         tags: list of tags to apply to generated py_library rules.
         entry_points: A dict of entry points to add py_binary rules for.
@@ -138,6 +165,10 @@ def generate_whl_library_build_bazel(
     srcs_exclude = []
     data_exclude = [] + data_exclude
     dependencies = sorted([normalize_name(d) for d in dependencies])
+    dependencies_by_platform = {
+        platform: sorted([normalize_name(d) for d in deps])
+        for platform, deps in dependencies_by_platform.items()
+    }
     tags = sorted(tags)
 
     for entry_point, entry_point_script_name in entry_points.items():
@@ -185,22 +216,48 @@ def generate_whl_library_build_bazel(
         for d in group_deps
     }
 
-    # Filter out deps which are within the group to avoid cycles
-    non_group_deps = [
+    dependencies = [
         d
         for d in dependencies
         if d not in group_deps
     ]
+    dependencies_by_platform = {
+        p: deps
+        for p, deps in dependencies_by_platform.items()
+        for deps in [[d for d in deps if d not in group_deps]]
+        if deps
+    }
 
-    lib_dependencies = [
-        "@%s%s//:%s" % (repo_prefix, normalize_name(d), PY_LIBRARY_PUBLIC_LABEL)
-        for d in non_group_deps
-    ]
+    for p in dependencies_by_platform:
+        if p.startswith("@"):
+            continue
 
-    whl_file_deps = [
-        "@%s%s//:%s" % (repo_prefix, normalize_name(d), WHEEL_FILE_PUBLIC_LABEL)
-        for d in non_group_deps
-    ]
+        os, _, cpu = p.partition("_")
+
+        additional_content.append(
+            """\
+config_setting(
+    name = "is_{os}_{cpu}",
+    constraint_values = [
+        "@platforms//cpu:{cpu}",
+        "@platforms//os:{os}",
+    ],
+    visibility = ["//visibility:private"],
+)
+""".format(os = os, cpu = cpu),
+        )
+
+    lib_dependencies = _render_list_and_select(
+        deps = dependencies,
+        deps_by_platform = dependencies_by_platform,
+        tmpl = "@{}{{}}//:{}".format(repo_prefix, PY_LIBRARY_PUBLIC_LABEL),
+    )
+
+    whl_file_deps = _render_list_and_select(
+        deps = dependencies,
+        deps_by_platform = dependencies_by_platform,
+        tmpl = "@{}{{}}//:{}".format(repo_prefix, WHEEL_FILE_PUBLIC_LABEL),
+    )
 
     # If this library is a member of a group, its public label aliases need to
     # point to the group implementation rule not the implementation rules. We
@@ -223,13 +280,13 @@ def generate_whl_library_build_bazel(
                 py_library_public_label = PY_LIBRARY_PUBLIC_LABEL,
                 py_library_impl_label = PY_LIBRARY_IMPL_LABEL,
                 py_library_actual_label = library_impl_label,
-                dependencies = repr(lib_dependencies),
+                dependencies = render.indent(lib_dependencies, " " * 4).lstrip(),
+                whl_file_deps = render.indent(whl_file_deps, " " * 4).lstrip(),
                 data_exclude = repr(_data_exclude),
                 whl_name = whl_name,
                 whl_file_public_label = WHEEL_FILE_PUBLIC_LABEL,
                 whl_file_impl_label = WHEEL_FILE_IMPL_LABEL,
                 whl_file_actual_label = whl_impl_label,
-                whl_file_deps = repr(whl_file_deps),
                 tags = repr(tags),
                 data_label = DATA_LABEL,
                 dist_info_label = DIST_INFO_LABEL,

python/pip_install/tools/wheel_installer/BUILD.bazel

@@ -13,6 +13,7 @@ py_library(
     deps = [
         requirement("installer"),
         requirement("pip"),
+        requirement("packaging"),
         requirement("setuptools"),
     ],
 )
@@ -47,6 +48,18 @@ py_test(
     ],
 )
 
+py_test(
+    name = "wheel_test",
+    size = "small",
+    srcs = [
+        "wheel_test.py",
+    ],
+    data = ["//examples/wheel:minimal_with_py_package"],
+    deps = [
+        ":lib",
+    ],
+)
+
 py_test(
     name = "wheel_installer_test",
     size = "small",

python/pip_install/tools/wheel_installer/arguments.py

@@ -15,7 +15,9 @@
 import argparse
 import json
 import pathlib
-from typing import Any
+from typing import Any, Dict, Set
+
+from python.pip_install.tools.wheel_installer import wheel
 
 
 def parser(**kwargs: Any) -> argparse.ArgumentParser:
@@ -39,6 +41,12 @@ def parser(**kwargs: Any) -> argparse.ArgumentParser:
         action="store",
         help="Extra arguments to pass down to pip.",
     )
+    parser.add_argument(
+        "--platform",
+        action="extend",
+        type=wheel.Platform.from_string,
+        help="Platforms to target dependencies. Can be used multiple times.",
+    )
     parser.add_argument(
         "--pip_data_exclude",
         action="store",
@@ -68,8 +76,9 @@ def parser(**kwargs: Any) -> argparse.ArgumentParser:
     return parser
 
 
-def deserialize_structured_args(args):
+def deserialize_structured_args(args: Dict[str, str]) -> Dict:
     """Deserialize structured arguments passed from the starlark rules.
+
     Args:
         args: dict of parsed command line arguments
     """
@@ -80,3 +89,18 @@ def deserialize_structured_args(args):
         else:
             args[arg_name] = []
     return args
+
+
+def get_platforms(args: argparse.Namespace) -> Set:
+    """Aggregate platforms into a single set.
+
+    Args:
+        args: dict of parsed command line arguments
+    """
+    platforms = set()
+    if args.platform is None:
+        return platforms
+
+    platforms.update(args.platform)
+
+    return platforms

python/pip_install/tools/wheel_installer/arguments_test.py

@@ -16,7 +16,7 @@
 import json
 import unittest
 
-from python.pip_install.tools.wheel_installer import arguments
+from python.pip_install.tools.wheel_installer import arguments, wheel
 
 
 class ArgumentsTestCase(unittest.TestCase):
@@ -52,6 +52,18 @@ def test_deserialize_structured_args(self) -> None:
         self.assertEqual(args["environment"], {"PIP_DO_SOMETHING": "True"})
         self.assertEqual(args["extra_pip_args"], [])
 
+    def test_platform_aggregation(self) -> None:
+        parser = arguments.parser()
+        args = parser.parse_args(
+            args=[
+                "--platform=host",
+                "--platform=linux_*",
+                "--platform=all",
+                "--requirement=foo",
+            ]
+        )
+        self.assertEqual(set(wheel.Platform.all()), arguments.get_platforms(args))
+
 
 if __name__ == "__main__":
     unittest.main()