From 3054c6ff5a1c96dc244e6e36a0e1e648421965cc Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 08:33:24 +0900
Subject: [PATCH 01/10] internal: support repo prefixes and config settings in
 alias rendering

Summary:
- introduce a whl_alias struct to make code more object oriented.
- make the interface of the function as small as possible.
- unify the bzlmod and legacy code paths.
- allow to specify arbitrary repo prefixes and config settings
  when generating aliases (split out from #1744).
---
 python/pip_install/pip_repository.bzl         |  10 +-
 python/private/bzlmod/pip.bzl                 |  18 ++-
 python/private/bzlmod/pip_repository.bzl      |  17 +-
 python/private/render_pkg_aliases.bzl         | 146 +++++++++---------
 .../render_pkg_aliases_test.bzl               |  94 ++++++-----
 5 files changed, 151 insertions(+), 134 deletions(-)

diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index 110ade19d7..5868ec523e 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -26,7 +26,7 @@ load("//python/private:envsubst.bzl", "envsubst")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:parse_whl_name.bzl", "parse_whl_name")
 load("//python/private:patch_whl.bzl", "patch_whl")
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")
 load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
 load("//python/private:toolchains_repo.bzl", "get_host_os_arch")
 load("//python/private:whl_target_platforms.bzl", "whl_target_platforms")
@@ -374,7 +374,13 @@ def _pip_repository_impl(rctx):
         config["experimental_target_platforms"] = rctx.attr.experimental_target_platforms
 
     macro_tmpl = "@%s//{}:{}" % rctx.attr.name
-    aliases = render_pkg_aliases(repo_name = rctx.attr.name, bzl_packages = bzl_packages)
+
+    aliases = render_pkg_aliases(
+        aliases = [
+            whl_alias(name = pkg, repo_prefix = rctx.attr.name + "_")
+            for pkg in bzl_packages or []
+        ],
+    )
     for path, contents in aliases.items():
         rctx.file(path, contents)
 
diff --git a/python/private/bzlmod/pip.bzl b/python/private/bzlmod/pip.bzl
index b4dbf2f1fe..767ecefb30 100644
--- a/python/private/bzlmod/pip.bzl
+++ b/python/private/bzlmod/pip.bzl
@@ -27,6 +27,7 @@ load(
 load("//python/pip_install:requirements_parser.bzl", parse_requirements = "parse")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:parse_whl_name.bzl", "parse_whl_name")
+load("//python/private:render_pkg_aliases.bzl", "whl_alias")
 load("//python/private:version_label.bzl", "version_label")
 load(":pip_repository.bzl", "pip_repository")
 
@@ -205,10 +206,14 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
             group_deps = group_deps,
         )
 
-        if whl_name not in whl_map[hub_name]:
-            whl_map[hub_name][whl_name] = {}
-
-        whl_map[hub_name][whl_name][_major_minor_version(pip_attr.python_version)] = pip_name + "_"
+        major_minor = _major_minor_version(pip_attr.python_version)
+        whl_map[hub_name].setdefault(whl_name, []).append(
+            whl_alias(
+                name = whl_name,
+                repo_prefix = pip_name,
+                version = major_minor,
+            ),
+        )
 
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories.
@@ -358,7 +363,10 @@ def _pip_impl(module_ctx):
         pip_repository(
             name = hub_name,
             repo_name = hub_name,
-            whl_map = whl_map,
+            whl_map = {
+                key: json.encode(value)
+                for key, value in whl_map.items()
+            },
             default_version = _major_minor_version(DEFAULT_PYTHON_VERSION),
         )
 
diff --git a/python/private/bzlmod/pip_repository.bzl b/python/private/bzlmod/pip_repository.bzl
index 8ea5ee7526..8fbd6f8d60 100644
--- a/python/private/bzlmod/pip_repository.bzl
+++ b/python/private/bzlmod/pip_repository.bzl
@@ -14,7 +14,7 @@
 
 ""
 
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")
 load("//python/private:text_util.bzl", "render")
 
 _BUILD_FILE_CONTENTS = """\
@@ -27,10 +27,12 @@ exports_files(["requirements.bzl"])
 def _pip_repository_impl(rctx):
     bzl_packages = rctx.attr.whl_map.keys()
     aliases = render_pkg_aliases(
-        repo_name = rctx.attr.repo_name,
-        rules_python = rctx.attr._template.workspace_name,
+        aliases = [
+            whl_alias(**v)
+            for key, values in rctx.attr.whl_map.items()
+            for v in json.decode(values)
+        ],
         default_version = rctx.attr.default_version,
-        whl_map = rctx.attr.whl_map,
     )
     for path, contents in aliases.items():
         rctx.file(path, contents)
@@ -71,9 +73,12 @@ setting.""",
         mandatory = True,
         doc = "The apparent name of the repo. This is needed because in bzlmod, the name attribute becomes the canonical name.",
     ),
-    "whl_map": attr.string_list_dict(
+    "whl_map": attr.string_dict(
         mandatory = True,
-        doc = "The wheel map where values are python versions",
+        doc = """\
+The wheel map where values are json.encoded strings of the whl_map constructed
+in the pip.parse tag class.
+""",
     ),
     "_template": attr.label(
         default = ":requirements.bzl.tmpl",
diff --git a/python/private/render_pkg_aliases.bzl b/python/private/render_pkg_aliases.bzl
index 02ba75bbc4..915392d5f7 100644
--- a/python/private/render_pkg_aliases.bzl
+++ b/python/private/render_pkg_aliases.bzl
@@ -16,9 +16,8 @@
 
 This is used in bzlmod and non-bzlmod setups."""
 
-load("//python/private:normalize_name.bzl", "normalize_name")
+load(":normalize_name.bzl", "normalize_name")
 load(":text_util.bzl", "render")
-load(":version_label.bzl", "version_label")
 
 NO_MATCH_ERROR_MESSAGE_TEMPLATE = """\
 No matching wheel for current configuration's Python version.
@@ -42,25 +41,18 @@ which has a "null" version value and will not match version constraints.
 def _render_whl_library_alias(
         *,
         name,
-        repo_name,
-        dep,
-        target,
         default_version,
-        versions,
-        rules_python):
-    """Render an alias for common targets
+        whl_repos):
+    """Render an alias for common targets."""
+    if len(whl_repos) == 1 and not whl_repos[0].version:
+        repo = whl_repos[0]
 
-    If the versions is passed, then the `rules_python` must be passed as well and
-    an alias with a select statement based on the python version is going to be
-    generated.
-    """
-    if versions == None:
         return render.alias(
             name = name,
-            actual = repr("@{repo_name}_{dep}//:{target}".format(
-                repo_name = repo_name,
-                dep = dep,
-                target = target,
+            actual = repr("@{repo_prefix}{dep}//:{name}".format(
+                repo_prefix = repo.repo_prefix,
+                dep = repo.name,
+                name = name,
             )),
         )
 
@@ -68,30 +60,21 @@ def _render_whl_library_alias(
     # statements  These select statements point to the different pip
     # whls that are based on a specific version of Python.
     selects = {}
-    for full_version in versions:
-        condition = "@@{rules_python}//python/config_settings:is_python_{full_python_version}".format(
-            rules_python = rules_python,
-            full_python_version = full_version,
-        )
-        actual = "@{repo_name}_{version}_{dep}//:{target}".format(
-            repo_name = repo_name,
-            version = version_label(full_version),
-            dep = dep,
-            target = target,
-        )
-        selects[condition] = actual
-
-    if default_version:
-        no_match_error = None
-        default_actual = "@{repo_name}_{version}_{dep}//:{target}".format(
-            repo_name = repo_name,
-            version = version_label(default_version),
-            dep = dep,
-            target = target,
+    no_match_error = "_NO_MATCH_ERROR"
+    default = None
+    for repo in sorted(whl_repos, key = lambda x: x.version):
+        actual = "@{repo_prefix}{dep}//:{name}".format(
+            repo_prefix = repo.repo_prefix,
+            dep = repo.name,
+            name = name,
         )
-        selects["//conditions:default"] = default_actual
-    else:
-        no_match_error = "_NO_MATCH_ERROR"
+        selects[repo.config_setting] = actual
+        if repo.version == default_version:
+            default = actual
+            no_match_error = None
+
+    if default:
+        selects["//conditions:default"] = default
 
     return render.alias(
         name = name,
@@ -101,22 +84,21 @@ def _render_whl_library_alias(
         ),
     )
 
-def _render_common_aliases(repo_name, name, versions = None, default_version = None, rules_python = None):
+def _render_common_aliases(*, name, whl_repos, default_version = None):
     lines = [
         """package(default_visibility = ["//visibility:public"])""",
     ]
 
-    if versions:
-        versions = sorted(versions)
+    versions = None
+    if whl_repos:
+        versions = sorted([v.version for v in whl_repos if v.version])
 
-    if not versions:
-        pass
-    elif default_version in versions:
+    if not versions or default_version in versions:
         pass
     else:
         error_msg = NO_MATCH_ERROR_MESSAGE_TEMPLATE.format(
             supported_versions = ", ".join(versions),
-            rules_python = rules_python,
+            rules_python = "rules_python",
         )
 
         lines.append("_NO_MATCH_ERROR = \"\"\"\\\n{error_msg}\"\"\"".format(
@@ -137,12 +119,8 @@ def _render_common_aliases(repo_name, name, versions = None, default_version = N
         [
             _render_whl_library_alias(
                 name = target,
-                repo_name = repo_name,
-                dep = name,
-                target = target,
-                versions = versions,
                 default_version = default_version,
-                rules_python = rules_python,
+                whl_repos = whl_repos,
             )
             for target in ["pkg", "whl", "data", "dist_info"]
         ],
@@ -150,7 +128,7 @@ def _render_common_aliases(repo_name, name, versions = None, default_version = N
 
     return "\n\n".join(lines)
 
-def render_pkg_aliases(*, repo_name, bzl_packages = None, whl_map = None, rules_python = None, default_version = None):
+def render_pkg_aliases(*, aliases, default_version = None):
     """Create alias declarations for each PyPI package.
 
     The aliases should be appended to the pip_repository BUILD.bazel file. These aliases
@@ -158,36 +136,58 @@ def render_pkg_aliases(*, repo_name, bzl_packages = None, whl_map = None, rules_
     when using bzlmod.
 
     Args:
-        repo_name: the repository name of the hub repository that is visible to the users that is
-            also used as the prefix for the spoke repo names (e.g. "pip", "pypi").
-        bzl_packages: the list of packages to setup, if not specified, whl_map.keys() will be used instead.
-        whl_map: the whl_map for generating Python version aware aliases.
+        aliases: the bzl_packages for generating Python version aware aliases.
         default_version: the default version to be used for the aliases.
-        rules_python: the name of the rules_python workspace.
 
     Returns:
         A dict of file paths and their contents.
     """
-    if not bzl_packages and whl_map:
-        bzl_packages = list(whl_map.keys())
-
     contents = {}
-    if not bzl_packages:
+    if not aliases:
         return contents
 
-    for name in bzl_packages:
-        versions = None
-        if whl_map != None:
-            versions = whl_map[name]
-        name = normalize_name(name)
+    whl_map = {}
+    for pkg in aliases:
+        whl_map.setdefault(pkg.name, []).append(pkg)
 
-        filename = "{}/BUILD.bazel".format(name)
-        contents[filename] = _render_common_aliases(
-            repo_name = repo_name,
+    return {
+        "{}/BUILD.bazel".format(name): _render_common_aliases(
             name = name,
-            versions = versions,
-            rules_python = rules_python,
+            whl_repos = whl_repos,
             default_version = default_version,
         ).strip()
+        for name, whl_repos in whl_map.items()
+    }
 
-    return contents
+def whl_alias(*, name, repo_prefix, version = None, config_setting = None):
+    """The bzl_packages value used by by the render_pkg_aliases function.
+
+    This contains the minimum amount of information required to generate correct
+    aliases in a hub repository.
+
+    Args:
+        name: str, the name of the package.
+        repo_prefix: str, the repo prefix of where to find the alias.
+        version: optional(str), the version of the python toolchain that this
+            whl alias is for. If not set, then non-version aware aliases will be
+            constructed. This is mainly used for better error messages when there
+            is no match found during a select.
+        config_setting: optional(Label or str), the config setting that we should use. Defaults
+            to "@rules_python//python/config_settings:is_python_{version}".
+
+    Returns:
+        a struct with the validated and parsed values.
+    """
+    if not repo_prefix:
+        fail("'repo_prefix' must be specified")
+
+    if version:
+        config_setting = config_setting or Label("//python/config_settings:is_python_" + version)
+        config_setting = str(config_setting)
+
+    return struct(
+        name = normalize_name(name),
+        repo_prefix = repo_prefix,
+        version = version,
+        config_setting = config_setting,
+    )
diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index 513c2783e9..0c93c98b62 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -15,14 +15,13 @@
 """render_pkg_aliases tests"""
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")  # buildifier: disable=bzl-visibility
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")  # buildifier: disable=bzl-visibility
 
 _tests = []
 
 def _test_empty(env):
     actual = render_pkg_aliases(
-        bzl_packages = None,
-        repo_name = "pypi",
+        aliases = None,
     )
 
     want = {}
@@ -33,8 +32,9 @@ _tests.append(_test_empty)
 
 def _test_legacy_aliases(env):
     actual = render_pkg_aliases(
-        bzl_packages = ["foo"],
-        repo_name = "pypi",
+        aliases = [
+            whl_alias(name = "foo", repo_prefix = "pypi_"),
+        ],
     )
 
     want = {
@@ -73,8 +73,10 @@ _tests.append(_test_legacy_aliases)
 
 def _test_all_legacy_aliases_are_created(env):
     actual = render_pkg_aliases(
-        bzl_packages = ["foo", "bar"],
-        repo_name = "pypi",
+        aliases = [
+            whl_alias(name = "bar", repo_prefix = "pypi_"),
+            whl_alias(name = "foo", repo_prefix = "pypi_"),
+        ],
     )
 
     want_files = ["bar/BUILD.bazel", "foo/BUILD.bazel"]
@@ -86,11 +88,9 @@ _tests.append(_test_all_legacy_aliases_are_created)
 def _test_bzlmod_aliases(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2"],
-        },
+        aliases = [
+            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
+        ],
     )
 
     want = {
@@ -106,7 +106,7 @@ alias(
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
             "//conditions:default": "@pypi_32_bar_baz//:pkg",
         },
     ),
@@ -116,7 +116,7 @@ alias(
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
             "//conditions:default": "@pypi_32_bar_baz//:whl",
         },
     ),
@@ -126,7 +126,7 @@ alias(
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
             "//conditions:default": "@pypi_32_bar_baz//:data",
         },
     ),
@@ -136,7 +136,7 @@ alias(
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
             "//conditions:default": "@pypi_32_bar_baz//:dist_info",
         },
     ),
@@ -150,11 +150,10 @@ _tests.append(_test_bzlmod_aliases)
 def _test_bzlmod_aliases_with_no_default_version(env):
     actual = render_pkg_aliases(
         default_version = None,
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2", "3.1"],
-        },
+        aliases = [
+            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
+            whl_alias(name = "bar-baz", version = "3.1", repo_prefix = "pypi_31_"),
+        ],
     )
 
     want_key = "bar_baz/BUILD.bazel"
@@ -189,8 +188,8 @@ alias(
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -200,8 +199,8 @@ alias(
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -211,8 +210,8 @@ alias(
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -222,8 +221,8 @@ alias(
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -244,11 +243,10 @@ def _test_bzlmod_aliases_for_non_root_modules(env):
         # non-root module, then we will have a no-match-error because the default_version
         # is not in the list of the versions in the whl_map.
         default_version = "3.3",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2", "3.1"],
-        },
+        aliases = [
+            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
+            whl_alias(name = "bar-baz", version = "3.1", repo_prefix = "pypi_31_"),
+        ],
     )
 
     want_key = "bar_baz/BUILD.bazel"
@@ -283,8 +281,8 @@ alias(
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -294,8 +292,8 @@ alias(
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -305,8 +303,8 @@ alias(
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -316,8 +314,8 @@ alias(
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -331,12 +329,12 @@ _tests.append(_test_bzlmod_aliases_for_non_root_modules)
 def _test_bzlmod_aliases_are_created_for_all_wheels(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar": ["3.1", "3.2"],
-            "foo": ["3.1", "3.2"],
-        },
+        aliases = [
+            whl_alias(name = "bar", version = "3.1", repo_prefix = "pypi_31_"),
+            whl_alias(name = "bar", version = "3.2", repo_prefix = "pypi_32_"),
+            whl_alias(name = "foo", version = "3.1", repo_prefix = "pypi_32_"),
+            whl_alias(name = "foo", version = "3.2", repo_prefix = "pypi_31_"),
+        ],
     )
 
     want_files = [

From cf84eb8e22e21fd7f84fa7b10bdc3922fe83503b Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 13:22:10 +0900
Subject: [PATCH 02/10] fixup! internal: support repo prefixes and config
 settings in alias rendering

---
 python/private/bzlmod/pip.bzl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/private/bzlmod/pip.bzl b/python/private/bzlmod/pip.bzl
index 767ecefb30..740a03d7fe 100644
--- a/python/private/bzlmod/pip.bzl
+++ b/python/private/bzlmod/pip.bzl
@@ -210,7 +210,7 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
         whl_map[hub_name].setdefault(whl_name, []).append(
             whl_alias(
                 name = whl_name,
-                repo_prefix = pip_name,
+                repo_prefix = pip_name + "_",
                 version = major_minor,
             ),
         )

From 1d8fb10cdb04034332b2b2b0b3aeca2357d9c079 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 13:42:16 +0900
Subject: [PATCH 03/10] Support the same test code on 6.4

---
 .../render_pkg_aliases/render_pkg_aliases_test.bzl | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index 0c93c98b62..f50475e53f 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -17,6 +17,16 @@
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")  # buildifier: disable=bzl-visibility
 
+def _normalize_labels(want):
+    # Do not modify the `want` on bazel 7+
+    if hasattr(native, "starlark_doc_extract"):
+        return want
+
+    return {
+        key: value.replace("\"@/", "\"@@/")
+        for key, value in want.items()
+    }
+
 _tests = []
 
 def _test_empty(env):
@@ -67,7 +77,7 @@ alias(
 )""",
     }
 
-    env.expect.that_dict(actual).contains_exactly(want)
+    env.expect.that_dict(actual).contains_exactly(_normalize_labels(want))
 
 _tests.append(_test_legacy_aliases)
 
@@ -143,7 +153,7 @@ alias(
 )""",
     }
 
-    env.expect.that_dict(actual).contains_exactly(want)
+    env.expect.that_dict(actual).contains_exactly(_normalize_labels(want))
 
 _tests.append(_test_bzlmod_aliases)
 

From 99b7907dafd1a80b925c164ced3311a78f059c25 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 14:01:40 +0900
Subject: [PATCH 04/10] Ensure that tests work the same under bzlmod and
 non-bzlmod

---
 .bazelversion                                 |  2 +-
 .../render_pkg_aliases_test.bzl               | 35 +++++++++----------
 2 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/.bazelversion b/.bazelversion
index 66ce77b7ea..19b860c187 100644
--- a/.bazelversion
+++ b/.bazelversion
@@ -1 +1 @@
-7.0.0
+6.4.0
diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index f50475e53f..d9e418ba0b 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -18,14 +18,14 @@ load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")  # buildifier: disable=bzl-visibility
 
 def _normalize_labels(want):
-    # Do not modify the `want` on bazel 7+
-    if hasattr(native, "starlark_doc_extract"):
+    if str(Label("//:invalid")).startswith("@@"):
+        # our expectations are already with double @
         return want
 
-    return {
-        key: value.replace("\"@/", "\"@@/")
-        for key, value in want.items()
-    }
+    if "@@" not in want:
+        fail("The expected string does not have '@@' labels, consider not using the function")
+
+    return want.replace("@@", "@")
 
 _tests = []
 
@@ -47,8 +47,8 @@ def _test_legacy_aliases(env):
         ],
     )
 
-    want = {
-        "foo/BUILD.bazel": """\
+    want_key = "foo/BUILD.bazel"
+    want_content = """\
 package(default_visibility = ["//visibility:public"])
 
 alias(
@@ -74,10 +74,9 @@ alias(
 alias(
     name = "dist_info",
     actual = "@pypi_foo//:dist_info",
-)""",
-    }
+)"""
 
-    env.expect.that_dict(actual).contains_exactly(_normalize_labels(want))
+    env.expect.that_dict(actual).contains_exactly({want_key: want_content})
 
 _tests.append(_test_legacy_aliases)
 
@@ -103,8 +102,8 @@ def _test_bzlmod_aliases(env):
         ],
     )
 
-    want = {
-        "bar_baz/BUILD.bazel": """\
+    want_key = "bar_baz/BUILD.bazel"
+    want_content = """\
 package(default_visibility = ["//visibility:public"])
 
 alias(
@@ -150,10 +149,10 @@ alias(
             "//conditions:default": "@pypi_32_bar_baz//:dist_info",
         },
     ),
-)""",
-    }
+)"""
 
-    env.expect.that_dict(actual).contains_exactly(_normalize_labels(want))
+    env.expect.that_collection(actual.keys()).contains_exactly([want_key])
+    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
 
 _tests.append(_test_bzlmod_aliases)
 
@@ -239,7 +238,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(want_content)
+    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
 
 _tests.append(_test_bzlmod_aliases_with_no_default_version)
 
@@ -332,7 +331,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(want_content)
+    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
 
 _tests.append(_test_bzlmod_aliases_for_non_root_modules)
 

From 7c8f47dec9b36b97662fe8abd49da90735dcf4e8 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 14:25:18 +0900
Subject: [PATCH 05/10] Update .bazelversion

---
 .bazelversion | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.bazelversion b/.bazelversion
index 19b860c187..66ce77b7ea 100644
--- a/.bazelversion
+++ b/.bazelversion
@@ -1 +1 @@
-6.4.0
+7.0.0

From ec848ec5a7f1a86a73da6fc51ed5b466d1fcdc17 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Wed, 14 Feb 2024 22:28:07 +0900
Subject: [PATCH 06/10] further cleanup and use dict as the main interface

---
 python/pip_install/pip_repository.bzl         |  6 +-
 python/private/bzlmod/pip.bzl                 |  7 +-
 python/private/bzlmod/pip_repository.bzl      |  7 +-
 python/private/render_pkg_aliases.bzl         | 62 +++++++---------
 .../render_pkg_aliases_test.bzl               | 70 +++++++++----------
 5 files changed, 69 insertions(+), 83 deletions(-)

diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index 5868ec523e..7b8160e956 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -376,10 +376,10 @@ def _pip_repository_impl(rctx):
     macro_tmpl = "@%s//{}:{}" % rctx.attr.name
 
     aliases = render_pkg_aliases(
-        aliases = [
-            whl_alias(name = pkg, repo_prefix = rctx.attr.name + "_")
+        aliases = {
+            pkg: [whl_alias(repo = rctx.attr.name + "_" + pkg)]
             for pkg in bzl_packages or []
-        ],
+        },
     )
     for path, contents in aliases.items():
         rctx.file(path, contents)
diff --git a/python/private/bzlmod/pip.bzl b/python/private/bzlmod/pip.bzl
index 740a03d7fe..a94dc28142 100644
--- a/python/private/bzlmod/pip.bzl
+++ b/python/private/bzlmod/pip.bzl
@@ -180,8 +180,9 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
         group_name = whl_group_mapping.get(whl_name)
         group_deps = requirement_cycles.get(group_name, [])
 
+        repo_name = "{}_{}".format(pip_name, whl_name)
         whl_library(
-            name = "%s_%s" % (pip_name, whl_name),
+            name = repo_name,
             requirement = requirement_line,
             repo = pip_name,
             repo_prefix = pip_name + "_",
@@ -209,9 +210,9 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
         major_minor = _major_minor_version(pip_attr.python_version)
         whl_map[hub_name].setdefault(whl_name, []).append(
             whl_alias(
-                name = whl_name,
-                repo_prefix = pip_name + "_",
+                repo = repo_name,
                 version = major_minor,
+                config_setting = Label("//python/config_settings:is_python_" + major_minor),
             ),
         )
 
diff --git a/python/private/bzlmod/pip_repository.bzl b/python/private/bzlmod/pip_repository.bzl
index 8fbd6f8d60..d96131dad7 100644
--- a/python/private/bzlmod/pip_repository.bzl
+++ b/python/private/bzlmod/pip_repository.bzl
@@ -27,11 +27,10 @@ exports_files(["requirements.bzl"])
 def _pip_repository_impl(rctx):
     bzl_packages = rctx.attr.whl_map.keys()
     aliases = render_pkg_aliases(
-        aliases = [
-            whl_alias(**v)
+        aliases = {
+            key: [whl_alias(**v) for v in json.decode(values)]
             for key, values in rctx.attr.whl_map.items()
-            for v in json.decode(values)
-        ],
+        },
         default_version = rctx.attr.default_version,
     )
     for path, contents in aliases.items():
diff --git a/python/private/render_pkg_aliases.bzl b/python/private/render_pkg_aliases.bzl
index 915392d5f7..e38f13321d 100644
--- a/python/private/render_pkg_aliases.bzl
+++ b/python/private/render_pkg_aliases.bzl
@@ -42,18 +42,13 @@ def _render_whl_library_alias(
         *,
         name,
         default_version,
-        whl_repos):
+        aliases):
     """Render an alias for common targets."""
-    if len(whl_repos) == 1 and not whl_repos[0].version:
-        repo = whl_repos[0]
-
+    if len(aliases) == 1 and not aliases[0].version:
+        alias = aliases[0]
         return render.alias(
             name = name,
-            actual = repr("@{repo_prefix}{dep}//:{name}".format(
-                repo_prefix = repo.repo_prefix,
-                dep = repo.name,
-                name = name,
-            )),
+            actual = repr("@{repo}//:{name}".format(repo = alias.repo, name = name)),
         )
 
     # Create the alias repositories which contains different select
@@ -62,14 +57,10 @@ def _render_whl_library_alias(
     selects = {}
     no_match_error = "_NO_MATCH_ERROR"
     default = None
-    for repo in sorted(whl_repos, key = lambda x: x.version):
-        actual = "@{repo_prefix}{dep}//:{name}".format(
-            repo_prefix = repo.repo_prefix,
-            dep = repo.name,
-            name = name,
-        )
-        selects[repo.config_setting] = actual
-        if repo.version == default_version:
+    for alias in sorted(aliases, key = lambda x: x.version):
+        actual = "@{repo}//:{name}".format(repo = alias.repo, name = name)
+        selects[alias.config_setting] = actual
+        if alias.version == default_version:
             default = actual
             no_match_error = None
 
@@ -84,14 +75,14 @@ def _render_whl_library_alias(
         ),
     )
 
-def _render_common_aliases(*, name, whl_repos, default_version = None):
+def _render_common_aliases(*, name, aliases, default_version = None):
     lines = [
         """package(default_visibility = ["//visibility:public"])""",
     ]
 
     versions = None
-    if whl_repos:
-        versions = sorted([v.version for v in whl_repos if v.version])
+    if aliases:
+        versions = sorted([v.version for v in aliases if v.version])
 
     if not versions or default_version in versions:
         pass
@@ -120,7 +111,7 @@ def _render_common_aliases(*, name, whl_repos, default_version = None):
             _render_whl_library_alias(
                 name = target,
                 default_version = default_version,
-                whl_repos = whl_repos,
+                aliases = aliases,
             )
             for target in ["pkg", "whl", "data", "dist_info"]
         ],
@@ -136,7 +127,8 @@ def render_pkg_aliases(*, aliases, default_version = None):
     when using bzlmod.
 
     Args:
-        aliases: the bzl_packages for generating Python version aware aliases.
+        aliases: dict, the keys are normalized distribution names and values are the
+            whl_alias instances.
         default_version: the default version to be used for the aliases.
 
     Returns:
@@ -145,29 +137,26 @@ def render_pkg_aliases(*, aliases, default_version = None):
     contents = {}
     if not aliases:
         return contents
-
-    whl_map = {}
-    for pkg in aliases:
-        whl_map.setdefault(pkg.name, []).append(pkg)
+    elif type(aliases) != type({}):
+        fail("The aliases need to be provided as a dict, got: {}".format(type(aliases)))
 
     return {
-        "{}/BUILD.bazel".format(name): _render_common_aliases(
-            name = name,
-            whl_repos = whl_repos,
+        "{}/BUILD.bazel".format(normalize_name(name)): _render_common_aliases(
+            name = normalize_name(name),
+            aliases = pkg_aliases,
             default_version = default_version,
         ).strip()
-        for name, whl_repos in whl_map.items()
+        for name, pkg_aliases in aliases.items()
     }
 
-def whl_alias(*, name, repo_prefix, version = None, config_setting = None):
+def whl_alias(*, repo, version = None, config_setting = None):
     """The bzl_packages value used by by the render_pkg_aliases function.
 
     This contains the minimum amount of information required to generate correct
     aliases in a hub repository.
 
     Args:
-        name: str, the name of the package.
-        repo_prefix: str, the repo prefix of where to find the alias.
+        repo: str, the repo of where to find the things to be aliased.
         version: optional(str), the version of the python toolchain that this
             whl alias is for. If not set, then non-version aware aliases will be
             constructed. This is mainly used for better error messages when there
@@ -178,16 +167,15 @@ def whl_alias(*, name, repo_prefix, version = None, config_setting = None):
     Returns:
         a struct with the validated and parsed values.
     """
-    if not repo_prefix:
-        fail("'repo_prefix' must be specified")
+    if not repo:
+        fail("'repo' must be specified")
 
     if version:
         config_setting = config_setting or Label("//python/config_settings:is_python_" + version)
         config_setting = str(config_setting)
 
     return struct(
-        name = normalize_name(name),
-        repo_prefix = repo_prefix,
+        repo = repo,
         version = version,
         config_setting = config_setting,
     )
diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index d9e418ba0b..807fc19155 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -42,9 +42,11 @@ _tests.append(_test_empty)
 
 def _test_legacy_aliases(env):
     actual = render_pkg_aliases(
-        aliases = [
-            whl_alias(name = "foo", repo_prefix = "pypi_"),
-        ],
+        aliases = {
+            "foo": [
+                whl_alias(repo = "pypi_foo"),
+            ],
+        },
     )
 
     want_key = "foo/BUILD.bazel"
@@ -80,26 +82,14 @@ alias(
 
 _tests.append(_test_legacy_aliases)
 
-def _test_all_legacy_aliases_are_created(env):
-    actual = render_pkg_aliases(
-        aliases = [
-            whl_alias(name = "bar", repo_prefix = "pypi_"),
-            whl_alias(name = "foo", repo_prefix = "pypi_"),
-        ],
-    )
-
-    want_files = ["bar/BUILD.bazel", "foo/BUILD.bazel"]
-
-    env.expect.that_dict(actual).keys().contains_exactly(want_files)
-
-_tests.append(_test_all_legacy_aliases_are_created)
-
 def _test_bzlmod_aliases(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        aliases = [
-            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
-        ],
+        aliases = {
+            "bar-baz": [
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+            ],
+        },
     )
 
     want_key = "bar_baz/BUILD.bazel"
@@ -159,10 +149,12 @@ _tests.append(_test_bzlmod_aliases)
 def _test_bzlmod_aliases_with_no_default_version(env):
     actual = render_pkg_aliases(
         default_version = None,
-        aliases = [
-            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
-            whl_alias(name = "bar-baz", version = "3.1", repo_prefix = "pypi_31_"),
-        ],
+        aliases = {
+            "bar-baz": [
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+                whl_alias(version = "3.1", repo = "pypi_31_bar_baz"),
+            ],
+        },
     )
 
     want_key = "bar_baz/BUILD.bazel"
@@ -252,10 +244,12 @@ def _test_bzlmod_aliases_for_non_root_modules(env):
         # non-root module, then we will have a no-match-error because the default_version
         # is not in the list of the versions in the whl_map.
         default_version = "3.3",
-        aliases = [
-            whl_alias(name = "bar-baz", version = "3.2", repo_prefix = "pypi_32_"),
-            whl_alias(name = "bar-baz", version = "3.1", repo_prefix = "pypi_31_"),
-        ],
+        aliases = {
+            "bar-baz": [
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+                whl_alias(version = "3.1", repo = "pypi_31_bar_baz"),
+            ],
+        },
     )
 
     want_key = "bar_baz/BUILD.bazel"
@@ -335,15 +329,19 @@ alias(
 
 _tests.append(_test_bzlmod_aliases_for_non_root_modules)
 
-def _test_bzlmod_aliases_are_created_for_all_wheels(env):
+def _test_aliases_are_created_for_all_wheels(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        aliases = [
-            whl_alias(name = "bar", version = "3.1", repo_prefix = "pypi_31_"),
-            whl_alias(name = "bar", version = "3.2", repo_prefix = "pypi_32_"),
-            whl_alias(name = "foo", version = "3.1", repo_prefix = "pypi_32_"),
-            whl_alias(name = "foo", version = "3.2", repo_prefix = "pypi_31_"),
-        ],
+        aliases = {
+            "bar": [
+                whl_alias(version = "3.1", repo = "pypi_31_bar"),
+                whl_alias(version = "3.2", repo = "pypi_32_bar"),
+            ],
+            "foo": [
+                whl_alias(version = "3.1", repo = "pypi_32_foo"),
+                whl_alias(version = "3.2", repo = "pypi_31_foo"),
+            ],
+        },
     )
 
     want_files = [
@@ -353,7 +351,7 @@ def _test_bzlmod_aliases_are_created_for_all_wheels(env):
 
     env.expect.that_dict(actual).keys().contains_exactly(want_files)
 
-_tests.append(_test_bzlmod_aliases_are_created_for_all_wheels)
+_tests.append(_test_aliases_are_created_for_all_wheels)
 
 def render_pkg_aliases_test_suite(name):
     """Create the test suite.

From 4ce9f6bfd177f2e182b85e12fc01446c882d3f2f Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Thu, 15 Feb 2024 13:25:49 +0900
Subject: [PATCH 07/10] Update python/private/bzlmod/pip.bzl

Co-authored-by: Richard Levasseur <richardlev@gmail.com>
---
 python/private/bzlmod/pip.bzl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/private/bzlmod/pip.bzl b/python/private/bzlmod/pip.bzl
index a94dc28142..a017089803 100644
--- a/python/private/bzlmod/pip.bzl
+++ b/python/private/bzlmod/pip.bzl
@@ -212,6 +212,7 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides):
             whl_alias(
                 repo = repo_name,
                 version = major_minor,
+                # Call Label() to canonicalize because its used in a different context
                 config_setting = Label("//python/config_settings:is_python_" + major_minor),
             ),
         )

From f0ecdf28ed43d50ae4bdbf8e60facb8558deeb51 Mon Sep 17 00:00:00 2001
From: aignas <240938+aignas@users.noreply.github.com>
Date: Thu, 15 Feb 2024 13:38:53 +0900
Subject: [PATCH 08/10] comment: explain the test expectation normalization
 function

---
 .../render_pkg_aliases_test.bzl               | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index 807fc19155..cce5477977 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -15,10 +15,19 @@
 """render_pkg_aliases tests"""
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
 load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")  # buildifier: disable=bzl-visibility
 
-def _normalize_labels(want):
-    if str(Label("//:invalid")).startswith("@@"):
+def _normalize_label_strings(want):
+    """normalize expected strings.
+
+    This function ensures that the desired `render_pkg_aliases` outputs are normalized from
+    `bzlmod` to `WORKSPACE` values so that we don't have to have to sets of expected strings.
+    The main difference is that under `bzlmod` the `str(Label("//my_label"))` results in
+    `"@@//my_label"` whereas under `non-bzlmod` we have `"@//my_label"`. This function does
+    `string.replace("@@", "@")` to normalize the strings.
+    """
+    if BZLMOD_ENABLED:
         # our expectations are already with double @
         return want
 
@@ -142,7 +151,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
+    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
 
 _tests.append(_test_bzlmod_aliases)
 
@@ -230,7 +239,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
+    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
 
 _tests.append(_test_bzlmod_aliases_with_no_default_version)
 
@@ -325,7 +334,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(_normalize_labels(want_content))
+    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
 
 _tests.append(_test_bzlmod_aliases_for_non_root_modules)
 

From d6256279ac77d635c1fef2f693b485c51e5b813a Mon Sep 17 00:00:00 2001
From: aignas <240938+aignas@users.noreply.github.com>
Date: Thu, 15 Feb 2024 13:44:26 +0900
Subject: [PATCH 09/10] improve tests and better explain the test helper

---
 .../render_pkg_aliases_test.bzl               | 30 ++++++++++++-------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index cce5477977..e1597238c2 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -21,11 +21,16 @@ load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias
 def _normalize_label_strings(want):
     """normalize expected strings.
 
-    This function ensures that the desired `render_pkg_aliases` outputs are normalized from
-    `bzlmod` to `WORKSPACE` values so that we don't have to have to sets of expected strings.
-    The main difference is that under `bzlmod` the `str(Label("//my_label"))` results in
-    `"@@//my_label"` whereas under `non-bzlmod` we have `"@//my_label"`. This function does
+    This function ensures that the desired `render_pkg_aliases` outputs are
+    normalized from `bzlmod` to `WORKSPACE` values so that we don't have to
+    have to sets of expected strings. The main difference is that under
+    `bzlmod` the `str(Label("//my_label"))` results in `"@@//my_label"` whereas
+    under `non-bzlmod` we have `"@//my_label"`. This function does
     `string.replace("@@", "@")` to normalize the strings.
+
+    NOTE, in tests, we should only use keep `@@` usage in expectation values
+    for the test cases where the whl_alias has the `config_setting` constructed
+    from a `Label` instance.
     """
     if BZLMOD_ENABLED:
         # our expectations are already with double @
@@ -96,7 +101,7 @@ def _test_bzlmod_aliases(env):
         default_version = "3.2",
         aliases = {
             "bar-baz": [
-                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz", config_setting = "//:my_config_setting"),
             ],
         },
     )
@@ -114,7 +119,7 @@ alias(
     name = "pkg",
     actual = select(
         {
-            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "//:my_config_setting": "@pypi_32_bar_baz//:pkg",
             "//conditions:default": "@pypi_32_bar_baz//:pkg",
         },
     ),
@@ -124,7 +129,7 @@ alias(
     name = "whl",
     actual = select(
         {
-            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "//:my_config_setting": "@pypi_32_bar_baz//:whl",
             "//conditions:default": "@pypi_32_bar_baz//:whl",
         },
     ),
@@ -134,7 +139,7 @@ alias(
     name = "data",
     actual = select(
         {
-            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "//:my_config_setting": "@pypi_32_bar_baz//:data",
             "//conditions:default": "@pypi_32_bar_baz//:data",
         },
     ),
@@ -144,7 +149,7 @@ alias(
     name = "dist_info",
     actual = select(
         {
-            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "//:my_config_setting": "@pypi_32_bar_baz//:dist_info",
             "//conditions:default": "@pypi_32_bar_baz//:dist_info",
         },
     ),
@@ -160,7 +165,12 @@ def _test_bzlmod_aliases_with_no_default_version(env):
         default_version = None,
         aliases = {
             "bar-baz": [
-                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+                whl_alias(
+                    version = "3.2",
+                    repo = "pypi_32_bar_baz",
+                    # pass the label to ensure that it gets converted to string
+                    config_setting = Label("//python/config_settings:is_python_3.2"),
+                ),
                 whl_alias(version = "3.1", repo = "pypi_31_bar_baz"),
             ],
         },

From 598e864cee079085c3b745959f21f1d025b4f0db Mon Sep 17 00:00:00 2001
From: aignas <240938+aignas@users.noreply.github.com>
Date: Thu, 15 Feb 2024 13:49:08 +0900
Subject: [PATCH 10/10] fixup! improve tests and better explain the test helper

---
 .../render_pkg_aliases/render_pkg_aliases_test.bzl        | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index e1597238c2..c61e5ef9b6 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -32,13 +32,13 @@ def _normalize_label_strings(want):
     for the test cases where the whl_alias has the `config_setting` constructed
     from a `Label` instance.
     """
+    if "@@" not in want:
+        fail("The expected string does not have '@@' labels, consider not using the function")
+
     if BZLMOD_ENABLED:
         # our expectations are already with double @
         return want
 
-    if "@@" not in want:
-        fail("The expected string does not have '@@' labels, consider not using the function")
-
     return want.replace("@@", "@")
 
 _tests = []
@@ -156,7 +156,7 @@ alias(
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
+    env.expect.that_str(actual[want_key]).equals(want_content)
 
 _tests.append(_test_bzlmod_aliases)