feat(whl_library): generate platform-specific dependency closures

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 done. This can be turned on by supplying
+  `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

@@ -102,6 +102,7 @@ pip.parse(
     python_version = "3.9",
     requirements_lock = "//:requirements_lock_3_9.txt",
     requirements_windows = "//:requirements_windows_3_9.txt",
+    target_platforms = ["all"],
     # These modifications were created above and we
     # are providing pip.parse with the label of the mod
     # and the name of the wheel.
@@ -125,6 +126,7 @@ pip.parse(
     python_version = "3.10",
     requirements_lock = "//:requirements_lock_3_10.txt",
     requirements_windows = "//:requirements_windows_3_10.txt",
+    target_platforms = ["all"],
     # These modifications were created above and we
     # are providing pip.parse with the label of the mod
     # and the name of the wheel.

python/pip_install/pip_repository.bzl

@@ -35,6 +35,19 @@ COMMAND_LINE_TOOLS_PATH_SLUG = "commandlinetools"
 
 _WHEEL_ENTRY_POINT_PREFIX = "rules_python_wheel_entry_point"
 
+_ALL_PLATFORMS = [
+    "linux_aarch64",
+    "linux_ppc64le",
+    "linux_s390x",
+    "linux_x86_64",
+    "linux_x86_32",
+    "osx_aarch64",
+    "osx_x86_64",
+    "windows_x86_64",
+    "windows_x86_32",
+    "windows_aarch64",
+]
+
 def _construct_pypath(rctx):
     """Helper function to construct a PYTHONPATH.
 
@@ -345,6 +358,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.target_platforms:
+        config["target_platforms"] = rctx.attr.target_platforms
 
     if rctx.attr.incompatible_generate_aliases:
         macro_tmpl = "@%s//{}:{}" % rctx.attr.name
@@ -513,6 +528,22 @@ python_interpreter. An example value: "@python3_x86_64-unknown-linux-gnu//:pytho
     "repo_prefix": attr.string(
         doc = """
 Prefix for the generated packages will be of the form `@<prefix><sanitized-package-name>//...`
+""",
+    ),
+    "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 when we use 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.
 """,
     ),
     # 600 is documented as default here: https://docs.bazel.build/versions/master/skylark/lib/repository_ctx.html#execute
@@ -669,6 +700,15 @@ See the example in rules_python/examples/pip_parse_vendored.
 )
 
 def _whl_library_impl(rctx):
+    target_platforms = []
+    if len(rctx.attr.target_platforms) == 1 and rctx.attr.target_platforms[0] == "all":
+        target_platforms = _ALL_PLATFORMS
+    elif rctx.attr.target_platforms:
+        target_platforms = rctx.attr.target_platforms
+        for p in target_platforms:
+            if p not in _ALL_PLATFORMS:
+                fail("target_platforms must be a subset of {} but got {}".format(["all"] + _ALL_PLATFORMS, target_platforms))
+
     python_interpreter = _resolve_python_interpreter(rctx)
     args = [
         python_interpreter,
@@ -713,7 +753,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 target_platforms],
         environment = environment,
         quiet = rctx.attr.quiet,
         timeout = rctx.attr.timeout,
@@ -749,6 +792,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 +859,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

@@ -17,6 +17,8 @@
 import pathlib
 from typing import Any
 
+from python.pip_install.tools.wheel_installer import wheel
+
 
 def parser(**kwargs: Any) -> argparse.ArgumentParser:
     """Create a parser for the wheel_installer tool."""
@@ -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="append",
+        type=wheel.Platform.from_string,
+        help="Platforms to target dependencies. Can be used multiple times.",
+    )
     parser.add_argument(
         "--pip_data_exclude",
         action="store",