feat(whl_library): set target_compatible_with based on whl name

This adds a pure starlark function to further parse the platform_tag
from the wheel and set the correct attributes for the
target generated by the whl_library rule. This should make errors
much more easy to understand when users try to build a docker image
using incompatible libraries.

This is currently hidden under an experimental flag described in the
changelog.

Author: aignas

CHANGELOG.md

@@ -97,6 +97,13 @@ Breaking changes:
 * (utils) Added a `pip_utils` struct with a `normalize_name` function to allow users
   to find out how `rules_python` would normalize a PyPI distribution name.
 
+* (whl_library) Added an `experimental_set_target_compatible_with` to `whl_library` and
+  `pip.parse` module extension which will configure bazel constraints for `py_library`
+  targets derived from a `whl` file via repository rules. In order to use it, you need to
+  specificy `experimental_set_target_compatible_with = True` to either `pip.parse` in your
+  `MODULE.bazel` or `install_deps` from the auto-generated `requirements.bzl` if you are
+  using `WORKSPACE`.
+
 ## [0.26.0] - 2023-10-06
 
 ### Changed

examples/bzlmod/MODULE.bazel

@@ -88,6 +88,8 @@ use_repo(pip, "whl_mods_hub")
 # Alternatively, `python_interpreter_target` can be used to directly specify
 # the Python interpreter to run to resolve dependencies.
 pip.parse(
+    # Experimental flags to test behaviour on CI, not part of example
+    experimental_set_target_compatible_with = True,
     hub_name = "pip",
     python_version = "3.9",
     requirements_lock = "//:requirements_lock_3_9.txt",
@@ -101,6 +103,8 @@ pip.parse(
     },
 )
 pip.parse(
+    # Experimental flags to test behaviour on CI, not part of example
+    experimental_set_target_compatible_with = True,
     hub_name = "pip",
     python_version = "3.10",
     requirements_lock = "//:requirements_lock_3_10.txt",

examples/bzlmod/other_module/MODULE.bazel

@@ -46,6 +46,8 @@ use_repo(
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 pip.parse(
+    # Experimental flags to test behaviour on CI, not part of example
+    experimental_set_target_compatible_with = True,
     hub_name = "other_module_pip",
     # NOTE: This version must be different than the root module's
     # default python version.

python/pip_install/pip_repository.bzl

@@ -668,6 +668,7 @@ def _whl_library_impl(rctx):
         ],
         entry_points = entry_points,
         annotation = None if not rctx.attr.annotation else struct(**json.decode(rctx.read(rctx.attr.annotation))),
+        set_target_compatible_with = rctx.attr.experimental_set_target_compatible_with,
     )
     rctx.file("BUILD.bazel", build_file_contents)
 
@@ -709,6 +710,11 @@ whl_library_attrs = {
         ),
         allow_files = True,
     ),
+    "experimental_set_target_compatible_with": attr.bool(
+        default = False,
+        doc = "A flag to set 'target_compatible_with' attribute for the py_library target. " +
+              "This is detected from the whl file name.",
+    ),
     "repo": attr.string(
         mandatory = True,
         doc = "Pointer to parent repo name. Used to make these rules rerun if the parent repo changes.",

python/pip_install/private/generate_whl_library_build_bazel.bzl

@@ -15,6 +15,7 @@
 """Generate the BUILD.bazel contents for a repo defined by a whl_library."""
 
 load("//python/private:normalize_name.bzl", "normalize_name")
+load("//python/private:parse_whl_name.bzl", "whl_target_compatible_with")
 
 _WHEEL_FILE_LABEL = "whl"
 _PY_LIBRARY_LABEL = "pkg"
@@ -77,6 +78,7 @@ py_library(
         ["site-packages/**/*"],
         exclude={data_exclude},
     ),
+    target_compatible_with = {target_compatible_with},
     # This makes this directory a top-level in the python import
     # search path for anything that depends on this.
     imports = ["site-packages"],
@@ -93,7 +95,8 @@ def generate_whl_library_build_bazel(
         data_exclude,
         tags,
         entry_points,
-        annotation = None):
+        annotation = None,
+        set_target_compatible_with = False):
     """Generate a BUILD file for an unzipped Wheel
 
     Args:
@@ -104,6 +107,8 @@ def generate_whl_library_build_bazel(
         tags: list of tags to apply to generated py_library rules.
         entry_points: A dict of entry points to add py_binary rules for.
         annotation: The annotation for the build file.
+        set_target_compatible_with: Set constraints for `py_library` based
+          on the whl name.
 
     Returns:
         A complete BUILD file as a string
@@ -178,6 +183,7 @@ def generate_whl_library_build_bazel(
                 entry_point_prefix = _WHEEL_ENTRY_POINT_PREFIX,
                 srcs_exclude = repr(srcs_exclude),
                 data = repr(data),
+                target_compatible_with = "None" if not set_target_compatible_with else repr(whl_target_compatible_with(whl_name)),
             ),
         ] + additional_content,
     )

python/private/bzlmod/pip.bzl

@@ -146,6 +146,7 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
             pip_data_exclude = pip_attr.pip_data_exclude,
             enable_implicit_namespace_pkgs = pip_attr.enable_implicit_namespace_pkgs,
             environment = pip_attr.environment,
+            experimental_set_target_compatible_with = pip_attr.experimental_set_target_compatible_with,
         )
 
         if whl_name not in whl_map[hub_name]:
@@ -307,6 +308,11 @@ def _pip_impl(module_ctx):
 
 def _pip_parse_ext_attrs():
     attrs = dict({
+        "experimental_set_target_compatible_with": attr.bool(
+            default = False,
+            doc = "A flag to set 'target_compatible_with' attribute for the py_library target. " +
+                  "This is detected from the whl file name.",
+        ),
         "hub_name": attr.string(
             mandatory = True,
             doc = """

python/private/parse_whl_name.bzl

@@ -16,6 +16,32 @@
 A starlark implementation of a Wheel filename parsing.
 """
 
+_LEGACY_ALIASES = {
+    "manylinux1_i686": "manylinux_2_5_i686",
+    "manylinux1_x86_64": "manylinux_2_5_x86_64",
+    "manylinux2010_i686": "manylinux_2_12_i686",
+    "manylinux2010_x86_64": "manylinux_2_12_x86_64",
+    "manylinux2014_aarch64": "manylinux_2_17_aarch64",
+    "manylinux2014_armv7l": "manylinux_2_17_armv7l",
+    "manylinux2014_i686": "manylinux_2_17_i686",
+    "manylinux2014_ppc64": "manylinux_2_17_ppc64",
+    "manylinux2014_ppc64le": "manylinux_2_17_ppc64le",
+    "manylinux2014_s390x": "manylinux_2_17_s390x",
+    "manylinux2014_x86_64": "manylinux_2_17_x86_64",
+}
+
+_ARCH = {
+    "aarch64": "aarch64",
+    "amd64": "x86_64",
+    "arm64": "aarch64",
+    "armv7l": "aarch32",
+    "i686": "x86_32",
+    "ppc64": "ppc",
+    "ppc64le": "ppc64le",
+    "s390x": "s390x",
+    "x86_64": "x86_64",
+}
+
 def parse_whl_name(file):
     """Parse whl file name into a struct of constituents.
 
@@ -70,3 +96,80 @@ def parse_whl_name(file):
         abi_tag = abi_tag,
         platform_tag = platform_tag,
     )
+
+def _convert_from_legacy(platform_tag):
+    return _LEGACY_ALIASES.get(platform_tag, platform_tag)
+
+def whl_target_compatible_with(file):
+    """Parse whl file and return compatibility list.
+
+    Args:
+        file (str): The file name of a wheel
+
+    Returns:
+        A list that can be put into target_compatible_with
+    """
+    parsed = parse_whl_name(file)
+
+    if parsed.platform_tag == "any" and parsed.abi_tag == "none":
+        return []
+
+    # TODO @aignas 2023-11-16: add ABI handling
+
+    platform, _, _ = parsed.platform_tag.partition(".")
+    platform = _convert_from_legacy(platform)
+
+    if platform.startswith("manylinux"):
+        _, _, tail = platform.partition("_")
+
+        _glibc_major, _, tail = tail.partition("_")  # Discard as this is currently unused
+        _glibc_minor, _, arch = tail.partition("_")  # Discard as this is currently unused
+
+        return [
+            "@platforms//cpu:" + _ARCH.get(arch, arch),
+            "@platforms//os:linux",
+        ]
+        # TODO @aignas 2023-11-16: figure out when this happens, perhaps it is when
+        # we build a wheel instead ourselves instead of downloading it from PyPI?
+
+    elif platform.startswith("linux_"):
+        _, _, arch = platform.partition("_")
+
+        return [
+            "@platforms//cpu:" + _ARCH.get(arch, arch),
+            "@platforms//os:linux",
+        ]
+
+    elif platform.startswith("macosx"):
+        _, _, tail = platform.partition("_")
+
+        _os_major, _, tail = tail.partition("_")  # Discard as this is currently unused
+        _os_minor, _, arch = tail.partition("_")  # Discard as this is currently unused
+
+        if arch.startswith("universal"):
+            return ["@platforms//os:osx"]
+        else:
+            return [
+                "@platforms//cpu:" + _ARCH.get(arch, arch),
+                "@platforms//os:osx",
+            ]
+    elif platform.startswith("win"):
+        if platform == "win32":
+            return [
+                "@platforms//cpu:x86_32",
+                "@platforms//os:windows",
+            ]
+        elif platform == "win64":
+            return [
+                "@platforms//cpu:x86_32",
+                "@platforms//os:windows",
+            ]
+
+        _, _, arch = platform.partition("_")
+
+        return [
+            "@platforms//cpu:" + _ARCH.get(arch, arch),
+            "@platforms//os:windows",
+        ]
+
+    fail("Could not parse platform values for a wheel platform: '{}'".format(parsed))

tests/pip_install/whl_library/generate_build_bazel_tests.bzl

@@ -55,6 +55,7 @@ py_library(
         ["site-packages/**/*"],
         exclude=["**/* *", "**/*.py", "**/*.pyc", "**/*.pyc.*", "**/*.dist-info/RECORD"],
     ),
+    target_compatible_with = None,
     # This makes this directory a top-level in the python import
     # search path for anything that depends on this.
     imports = ["site-packages"],
@@ -111,6 +112,7 @@ py_library(
         ["site-packages/**/*"],
         exclude=["**/* *", "**/*.py", "**/*.pyc", "**/*.pyc.*", "**/*.dist-info/RECORD", "data_exclude_all"],
     ),
+    target_compatible_with = None,
     # This makes this directory a top-level in the python import
     # search path for anything that depends on this.
     imports = ["site-packages"],
@@ -190,6 +192,7 @@ py_library(
         ["site-packages/**/*"],
         exclude=["**/* *", "**/*.py", "**/*.pyc", "**/*.pyc.*", "**/*.dist-info/RECORD"],
     ),
+    target_compatible_with = None,
     # This makes this directory a top-level in the python import
     # search path for anything that depends on this.
     imports = ["site-packages"],

tests/private/whl_target_compatible_with/BUILD.bazel

@@ -0,0 +1,3 @@
+load(":whl_target_compatible_with_tests.bzl", "whl_target_compatible_with_test_suite")
+
+whl_target_compatible_with_test_suite(name = "whl_target_compatible_with_tests")

tests/private/whl_target_compatible_with/whl_target_compatible_with_tests.bzl

@@ -0,0 +1,57 @@
+# 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.
+
+""
+
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("//python/private:parse_whl_name.bzl", "whl_target_compatible_with")  # buildifier: disable=bzl-visibility
+
+_tests = []
+
+def _test_compatible_with_all(env):
+    got = whl_target_compatible_with("foo-1.2.3-py3-none-any.whl")
+    env.expect.that_collection(got).contains_exactly([])
+
+_tests.append(_test_compatible_with_all)
+
+def _test_multiple_platforms(env):
+    got = whl_target_compatible_with("bar-3.2.1-py3-abi3-manylinux1_x86_64.manylinux2_x86_64.whl")
+    env.expect.that_collection(got).contains_exactly([
+        "@platforms//os:linux",
+        "@platforms//cpu:x86_64",
+    ])
+
+_tests.append(_test_multiple_platforms)
+
+def _test_real_numpy_wheel(env):
+    got = whl_target_compatible_with("numpy-1.26.1-pp39-pypy39_pp73-macosx_10_9_x86_64.whl")
+    env.expect.that_collection(got).contains_exactly([
+        "@platforms//os:osx",
+        "@platforms//cpu:x86_64",
+    ])
+
+_tests.append(_test_real_numpy_wheel)
+
+# TODO @aignas 2023-11-16: add handling for musllinux as simple linux for now.
+# TODO @aignas 2023-11-16: add macos universal testcase.
+# TODO @aignas 2023-11-16: add Windows testcases.
+# TODO @aignas 2023-11-16: add error handling when the wheel filename is something else.
+
+def whl_target_compatible_with_test_suite(name):
+    """Create the test suite.
+
+    Args:
+        name: the name of the test suite
+    """
+    test_suite(name = name, basic_tests = _tests)