refactor(pkg_aliases): create a macro for creating whl aliases (#2391)

This just cleans up the code and moves more logic from the
repository_rule
(i.e. generation of `BUILD.bazel` files) to loading time (macro
evaluation).
This makes the unit testing easier and I plan to also move the code that
is
generating config setting names from filenames to this new macro, but
wanted to
submit this PR to reduce the review chunks.

Summary:
- Add a new `pkg_aliases` macro.
- Move logic and tests for creating WORKSPACE aliases.
- Move logic and tests bzlmod aliases.
- Move logic and tests bzlmod aliases with groups.
- Add a test for extra alias creation.
- Use `whl_alias` in `pypi` extension integration tests.
- Improve the serialization of `whl_alias` for passing to the pypi hub
repo.

Related to #260, #2386, #2337, #2319 - hopefully cleaning the code up
will make
it easier to address those feature requests later.

---------

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

Author: aignas

examples/bzlmod/MODULE.bazel.lock

@@ -256,7 +256,6 @@ bzl_library(
     srcs = ["render_pkg_aliases.bzl"],
     deps = [
         ":generate_group_library_build_bazel_bzl",
-        ":labels_bzl",
         ":parse_whl_name_bzl",
         ":whl_target_platforms_bzl",
         "//python/private:normalize_name_bzl",

python/private/pypi/BUILD.bazel

@@ -571,6 +571,20 @@ You cannot use both the additive_build_content and additive_build_content_file a
         is_reproducible = is_reproducible,
     )
 
+def _alias_dict(a):
+    ret = {
+        "repo": a.repo,
+    }
+    if a.config_setting:
+        ret["config_setting"] = a.config_setting
+    if a.filename:
+        ret["filename"] = a.filename
+    if a.target_platforms:
+        ret["target_platforms"] = a.target_platforms
+    if a.version:
+        ret["version"] = a.version
+    return ret
+
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories.
 
@@ -651,8 +665,8 @@ def _pip_impl(module_ctx):
             repo_name = hub_name,
             extra_hub_aliases = mods.extra_aliases.get(hub_name, {}),
             whl_map = {
-                key: json.encode(value)
-                for key, value in whl_map.items()
+                key: json.encode([_alias_dict(a) for a in aliases])
+                for key, aliases in whl_map.items()
             },
             packages = mods.exposed_packages.get(hub_name, []),
             groups = mods.hub_group_map.get(hub_name),

python/private/pypi/extension.bzl

@@ -0,0 +1,134 @@
+# Copyright 2024 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.
+
+"""pkg_aliases is a macro to generate aliases for selecting the right wheel for the right target platform.
+
+This is used in bzlmod and non-bzlmod setups."""
+
+load("//python/private:text_util.bzl", "render")
+load(
+    ":labels.bzl",
+    "DATA_LABEL",
+    "DIST_INFO_LABEL",
+    "PY_LIBRARY_IMPL_LABEL",
+    "PY_LIBRARY_PUBLIC_LABEL",
+    "WHEEL_FILE_IMPL_LABEL",
+    "WHEEL_FILE_PUBLIC_LABEL",
+)
+
+_NO_MATCH_ERROR_TEMPLATE = """\
+No matching wheel for current configuration's Python version.
+
+The current build configuration's Python version doesn't match any of the Python
+wheels available for this wheel. This wheel supports the following Python
+configuration settings:
+    {config_settings}
+
+To determine the current configuration's Python version, run:
+    `bazel config <config id>` (shown further below)
+and look for
+    {rules_python}//python/config_settings:python_version
+
+If the value is missing, then the "default" Python version is being used,
+which has a "null" version value and will not match version constraints.
+"""
+
+def _no_match_error(actual):
+    if type(actual) != type({}):
+        return None
+
+    if "//conditions:default" in actual:
+        return None
+
+    return _NO_MATCH_ERROR_TEMPLATE.format(
+        config_settings = render.indent(
+            "\n".join(sorted(actual.keys())),
+        ).lstrip(),
+        rules_python = "rules_python",
+    )
+
+def pkg_aliases(
+        *,
+        name,
+        actual,
+        group_name = None,
+        extra_aliases = None,
+        native = native,
+        select = select):
+    """Create aliases for an actual package.
+
+    Args:
+        name: {type}`str` The name of the package.
+        actual: {type}`dict[Label, str] | str` The name of the repo the aliases point to, or a dict of select conditions to repo names for the aliases to point to
+            mapping to repositories.
+        group_name: {type}`str` The group name that the pkg belongs to.
+        extra_aliases: {type}`list[str]` The extra aliases to be created.
+        native: {type}`struct` used in unit tests.
+        select: {type}`select` used in unit tests.
+    """
+    native.alias(
+        name = name,
+        actual = ":" + PY_LIBRARY_PUBLIC_LABEL,
+    )
+
+    target_names = {
+        PY_LIBRARY_PUBLIC_LABEL: PY_LIBRARY_IMPL_LABEL if group_name else PY_LIBRARY_PUBLIC_LABEL,
+        WHEEL_FILE_PUBLIC_LABEL: WHEEL_FILE_IMPL_LABEL if group_name else WHEEL_FILE_PUBLIC_LABEL,
+        DATA_LABEL: DATA_LABEL,
+        DIST_INFO_LABEL: DIST_INFO_LABEL,
+    } | {
+        x: x
+        for x in extra_aliases or []
+    }
+    no_match_error = _no_match_error(actual)
+
+    for name, target_name in target_names.items():
+        if type(actual) == type(""):
+            _actual = "@{repo}//:{target_name}".format(
+                repo = actual,
+                target_name = name,
+            )
+        elif type(actual) == type({}):
+            _actual = select(
+                {
+                    config_setting: "@{repo}//:{target_name}".format(
+                        repo = repo,
+                        target_name = name,
+                    )
+                    for config_setting, repo in actual.items()
+                },
+                no_match_error = no_match_error,
+            )
+        else:
+            fail("The `actual` arg must be a dictionary or a string")
+
+        kwargs = {}
+        if target_name.startswith("_"):
+            kwargs["visibility"] = ["//_groups:__subpackages__"]
+
+        native.alias(
+            name = target_name,
+            actual = _actual,
+            **kwargs
+        )
+
+    if group_name:
+        native.alias(
+            name = PY_LIBRARY_PUBLIC_LABEL,
+            actual = "//_groups:{}_pkg".format(group_name),
+        )
+        native.alias(
+            name = WHEEL_FILE_PUBLIC_LABEL,
+            actual = "//_groups:{}_whl".format(group_name),
+        )

python/private/pypi/pkg_aliases.bzl

@@ -22,15 +22,6 @@ load(
     ":generate_group_library_build_bazel.bzl",
     "generate_group_library_build_bazel",
 )  # buildifier: disable=bzl-visibility
-load(
-    ":labels.bzl",
-    "DATA_LABEL",
-    "DIST_INFO_LABEL",
-    "PY_LIBRARY_IMPL_LABEL",
-    "PY_LIBRARY_PUBLIC_LABEL",
-    "WHEEL_FILE_IMPL_LABEL",
-    "WHEEL_FILE_PUBLIC_LABEL",
-)
 load(":parse_whl_name.bzl", "parse_whl_name")
 load(":whl_target_platforms.bzl", "whl_target_platforms")
 
@@ -70,117 +61,32 @@ If the value is missing, then the "default" Python version is being used,
 which has a "null" version value and will not match version constraints.
 """
 
-def _render_whl_library_alias(
-        *,
-        name,
-        aliases,
-        target_name,
-        **kwargs):
-    """Render an alias for common targets."""
-    if len(aliases) == 1 and not aliases[0].version:
-        alias = aliases[0]
-        return render.alias(
-            name = name,
-            actual = repr("@{repo}//:{name}".format(
-                repo = alias.repo,
-                name = target_name,
-            )),
-            **kwargs
-        )
-
-    # Create the alias repositories which contains different select
-    # statements  These select statements point to the different pip
-    # whls that are based on a specific version of Python.
-    selects = {}
-    no_match_error = "_NO_MATCH_ERROR"
-    for alias in sorted(aliases, key = lambda x: x.version):
-        actual = "@{repo}//:{name}".format(repo = alias.repo, name = target_name)
-        selects.setdefault(actual, []).append(alias.config_setting)
+def _repr_actual(aliases):
+    if len(aliases) == 1 and not aliases[0].version and not aliases[0].config_setting:
+        return repr(aliases[0].repo)
 
-    return render.alias(
-        name = name,
-        actual = render.select(
-            {
-                tuple(sorted(
-                    conditions,
-                    # Group `is_python` and other conditions for easier reading
-                    # when looking at the generated files.
-                    key = lambda condition: ("is_python" not in condition, condition),
-                )): target
-                for target, conditions in sorted(selects.items())
-            },
-            no_match_error = no_match_error,
-            # This key_repr is used to render selects.with_or keys
-            key_repr = lambda x: repr(x[0]) if len(x) == 1 else render.tuple(x),
-            name = "selects.with_or",
-        ),
-        **kwargs
-    )
+    actual = {}
+    for alias in aliases:
+        actual[alias.config_setting or ("//_config:is_python_" + alias.version)] = alias.repo
+    return render.indent(render.dict(actual)).lstrip()
 
 def _render_common_aliases(*, name, aliases, extra_aliases = [], group_name = None):
-    lines = [
-        """load("@bazel_skylib//lib:selects.bzl", "selects")""",
-        """package(default_visibility = ["//visibility:public"])""",
-    ]
-
-    config_settings = None
-    if aliases:
-        config_settings = sorted([v.config_setting for v in aliases if v.config_setting])
-
-    if config_settings:
-        error_msg = NO_MATCH_ERROR_MESSAGE_TEMPLATE_V2.format(
-            config_settings = render.indent(
-                "\n".join(config_settings),
-            ).lstrip(),
-            rules_python = "rules_python",
-        )
+    return """\
+load("@rules_python//python/private/pypi:pkg_aliases.bzl", "pkg_aliases")
 
-        lines.append("_NO_MATCH_ERROR = \"\"\"\\\n{error_msg}\"\"\"".format(
-            error_msg = error_msg,
-        ))
+package(default_visibility = ["//visibility:public"])
 
-    lines.append(
-        render.alias(
-            name = name,
-            actual = repr(":pkg"),
-        ),
-    )
-    lines.extend(
-        [
-            _render_whl_library_alias(
-                name = name,
-                aliases = aliases,
-                target_name = target_name,
-                visibility = ["//_groups:__subpackages__"] if name.startswith("_") else None,
-            )
-            for target_name, name in (
-                {
-                    PY_LIBRARY_PUBLIC_LABEL: PY_LIBRARY_IMPL_LABEL if group_name else PY_LIBRARY_PUBLIC_LABEL,
-                    WHEEL_FILE_PUBLIC_LABEL: WHEEL_FILE_IMPL_LABEL if group_name else WHEEL_FILE_PUBLIC_LABEL,
-                    DATA_LABEL: DATA_LABEL,
-                    DIST_INFO_LABEL: DIST_INFO_LABEL,
-                } | {
-                    x: x
-                    for x in extra_aliases
-                }
-            ).items()
-        ],
+pkg_aliases(
+    name = "{name}",
+    actual = {actual},
+    group_name = {group_name},
+    extra_aliases = {extra_aliases},
+)""".format(
+        name = name,
+        actual = _repr_actual(aliases),
+        group_name = repr(group_name),
+        extra_aliases = repr(extra_aliases),
     )
-    if group_name:
-        lines.extend(
-            [
-                render.alias(
-                    name = "pkg",
-                    actual = repr("//_groups:{}_pkg".format(group_name)),
-                ),
-                render.alias(
-                    name = "whl",
-                    actual = repr("//_groups:{}_whl".format(group_name)),
-                ),
-            ],
-        )
-
-    return "\n\n".join(lines)
 
 def render_pkg_aliases(*, aliases, requirement_cycles = None, extra_hub_aliases = {}):
     """Create alias declarations for each PyPI package.
@@ -222,7 +128,7 @@ def render_pkg_aliases(*, aliases, requirement_cycles = None, extra_hub_aliases
         "{}/BUILD.bazel".format(normalize_name(name)): _render_common_aliases(
             name = normalize_name(name),
             aliases = pkg_aliases,
-            extra_aliases = extra_hub_aliases.get(name, []),
+            extra_aliases = extra_hub_aliases.get(normalize_name(name), []),
             group_name = whl_group_mapping.get(normalize_name(name)),
         ).strip()
         for name, pkg_aliases in aliases.items()
@@ -256,21 +162,17 @@ def whl_alias(*, repo, version = None, config_setting = None, filename = None, t
     if not repo:
         fail("'repo' must be specified")
 
-    if version:
-        config_setting = config_setting or ("//_config:is_python_" + version)
-        config_setting = str(config_setting)
-
     if target_platforms:
         for p in target_platforms:
             if not p.startswith("cp"):
                 fail("target_platform should start with 'cp' denoting the python version, got: " + p)
 
     return struct(
-        repo = repo,
-        version = version,
         config_setting = config_setting,
         filename = filename,
+        repo = repo,
         target_platforms = target_platforms,
+        version = version,
     )
 
 def render_multiplatform_pkg_aliases(*, aliases, **kwargs):