internal: support repo prefixes and config settings in alias rendering

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 = {
+            pkg: [whl_alias(repo = rctx.attr.name + "_" + pkg)]
+            for pkg in bzl_packages or []
+        },
+    )
     for path, contents in aliases.items():
         rctx.file(path, contents)
 

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")
 
@@ -179,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 + "_",
@@ -205,10 +207,15 @@ 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(
+                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),
+            ),
+        )
 
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories.
@@ -358,7 +365,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),
         )
 

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,11 @@ 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 = {
+            key: [whl_alias(**v) for v in json.decode(values)]
+            for key, values in rctx.attr.whl_map.items()
+        },
         default_version = rctx.attr.default_version,
-        whl_map = rctx.attr.whl_map,
     )
     for path, contents in aliases.items():
         rctx.file(path, contents)
@@ -71,9 +72,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",

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,56 +41,31 @@ 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
-
-    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:
+        aliases):
+    """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}_{dep}//:{target}".format(
-                repo_name = repo_name,
-                dep = dep,
-                target = target,
-            )),
+            actual = repr("@{repo}//:{name}".format(repo = alias.repo, name = name)),
         )
 
     # 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 = {}
-    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,
-        )
-        selects["//conditions:default"] = default_actual
-    else:
-        no_match_error = "_NO_MATCH_ERROR"
+    no_match_error = "_NO_MATCH_ERROR"
+    default = None
+    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
+
+    if default:
+        selects["//conditions:default"] = default
 
     return render.alias(
         name = name,
@@ -101,22 +75,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, aliases, default_version = None):
     lines = [
         """package(default_visibility = ["//visibility:public"])""",
     ]
 
-    if versions:
-        versions = sorted(versions)
+    versions = None
+    if aliases:
+        versions = sorted([v.version for v in aliases 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,57 +110,72 @@ 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,
+                aliases = aliases,
             )
             for target in ["pkg", "whl", "data", "dist_info"]
         ],
     )
 
     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
     allow users to use requirement() without needed a corresponding `use_repo()` for each dep
     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: 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.
-        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
+    elif type(aliases) != type({}):
+        fail("The aliases need to be provided as a dict, got: {}".format(type(aliases)))
 
-    for name in bzl_packages:
-        versions = None
-        if whl_map != None:
-            versions = whl_map[name]
-        name = normalize_name(name)
-
-        filename = "{}/BUILD.bazel".format(name)
-        contents[filename] = _render_common_aliases(
-            repo_name = repo_name,
-            name = name,
-            versions = versions,
-            rules_python = rules_python,
+    return {
+        "{}/BUILD.bazel".format(normalize_name(name)): _render_common_aliases(
+            name = normalize_name(name),
+            aliases = pkg_aliases,
             default_version = default_version,
         ).strip()
+        for name, pkg_aliases in aliases.items()
+    }
 
-    return contents
+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:
+        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
+            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:
+        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(
+        repo = repo,
+        version = version,
+        config_setting = config_setting,
+    )