feat: allow toolchain without host compatible variant

Author: rickeylev

python/private/python.bzl

@@ -21,7 +21,13 @@ load(":full_version.bzl", "full_version")
 load(":python_register_toolchains.bzl", "python_register_toolchains")
 load(":pythons_hub.bzl", "hub_repo")
 load(":repo_utils.bzl", "repo_utils")
-load(":toolchains_repo.bzl", "host_toolchain", "multi_toolchain_aliases", "sorted_host_platforms")
+load(
+    ":toolchains_repo.bzl",
+    "host_toolchain",
+    "multi_toolchain_aliases",
+    "sorted_host_platform_names",
+    "sorted_host_platforms",
+)
 load(":util.bzl", "IS_BAZEL_6_4_OR_HIGHER")
 load(":version.bzl", "version")
 
@@ -267,6 +273,8 @@ def parse_modules(*, module_ctx, _fail = fail):
 def _python_impl(module_ctx):
     py = parse_modules(module_ctx = module_ctx)
 
+    all_host_compatible_impls = []
+
     # list of structs; see inline struct call within the loop below.
     toolchain_impls = []
 
@@ -319,22 +327,96 @@ def _python_impl(module_ctx):
             ))
             if _is_compatible_with_host(module_ctx, platform_info):
                 host_platforms[platform_name] = platform_info
-
-        host_platforms = sorted_host_platforms(host_platforms)
-        host_toolchain(
-            name = toolchain_info.name + "_host",
-            # NOTE: Order matters. The first found to be compatible is (usually) used.
-            platforms = host_platforms.keys(),
-            os_names = {
-                str(i): platform_info.os_name
-                for i, platform_info in enumerate(host_platforms.values())
-            },
+                host_compat_entry = struct(
+                    full_python_version = full_python_version,
+                    platform = platform_info,
+                    platform_name = platform_name,
+                    impl_repo_name = repo_name,
+                )
+                all_host_compatible_impls.setdefault(full_python_version, []).append(
+                    host_compat_entry,
+                )
+                all_host_compatible_impls.setdefault(
+                    toolchain_info.python_version,
+                    [],
+                ).append(host_compat_entry)
+
+        host_repo_name = toolchain_info.name + "_host"
+        if not host_platforms:
+            needed_host_repos[toolchain_info.python_version] = struct(
+                compatible_version = toolchain_info.python_version,
+                full_python_version = full_python_version,
+            )
+        else:
+            host_platforms = sorted_host_platforms(host_platforms)
+            host_toolchain(
+                name = host_repo_name,
+                # NOTE: Order matters. The first found to be compatible is (usually) used.
+                platforms = host_platforms.keys(),
+                os_names = {
+                    str(i): platform_info.os_name
+                    for i, platform_info in enumerate(host_platforms.values())
+                },
             arch_names = {
-                str(i): platform_info.arch
-                for i, platform_info in enumerate(host_platforms.values())
-            },
-            python_version = full_python_version,
-        )
+                    str(i): platform_info.arch
+                    for i, platform_info in enumerate(host_platforms.values())
+                },
+                python_version = full_python_version,
+            )
+
+    """
+    We want to define e.g. python_3_13_host backed by e.g. python_3_13 (full
+    version 3.13.3)
+    but, we didn't see a host-compatible option when we initial did so.
+    So search from 3.13.3 down to 3.13.0, looking for compatible options.
+    """
+
+    def vt(s):
+        return tuple([int(x) for x in s.split(".")])
+
+    if needed_host_repos:
+        for key, entries in all_host_compatible_impls.items():
+            all_host_compatible_impls[key] = sorted(
+                entries,
+                reverse = True,
+                key = lambda e: vt(e.full_python_version),
+            )
+
+    for host_repo_name, info in needed_host_repos.items():
+        choices = []
+        for entry in all_host_compatible_impls[info.compatible_version]:
+            # todo: numeric version comparison
+            if vt(entry.full_python_version) <= vt(info.full_python_version):
+                choices.append(entry)
+        if choices:
+            platforms_keys = [
+                # We have to prepend the offset because the same platform
+                # name might occur accross different versions
+                "{}_{}".format(i, entry.platform_name)
+                for i, entry in enumerate(choices)
+            ]
+            platform_keys = sorted_host_platforms_names(platform_keys)
+
+            # AH no, this won't quite work.
+            # Multiple versions will have the same platform_name string.
+            # Thus when it builds the platform_map, some will get clobbered
+            # Maybe just throw an offset prefix onto it for uniqueness?
+            host_toolchain(
+                name = host_repo_name,
+                platforms = platform_keys,
+                impl_repo_suffixes = {
+                    # Internally, it prepends its own name arg, so we must remove
+                    # the common prefix
+                    str(i): entry.impl_repo_name.removeprefix(host_repo_name)
+                    for i, entry in enumerate(choices)
+                },
+                os_names = {str(i): entry.os_name for entry in enumerate(choices)},
+                archs = {str(i): entry.arch for entry in enumerate(choices)},
+                python_versions = {str(i): entry.full_python_version for entry in enumerate(choices)},
+            )
+        else:
+            # todo: figure out what to do. Define nothing, if we can.
+            fail("No host-compatible found")
 
     # List of the base names ("python_3_10") for the toolchain repos
     base_toolchain_repo_names = []

python/private/toolchains_repo.bzl

@@ -313,7 +313,7 @@ def _host_toolchain_impl(rctx):
     rctx.file("BUILD.bazel", _HOST_TOOLCHAIN_BUILD_CONTENT)
 
     os_name = repo_utils.get_platforms_os_name(rctx)
-    host_platform = _get_host_platform(
+    host_platform = _get_host_impl_repo_suffix(
         rctx = rctx,
         logger = repo_utils.logger(rctx),
         python_version = rctx.attr.python_version,
@@ -386,7 +386,27 @@ Creates a repository with a shorter name meant to be used in the repository_ctx,
 which needs to have `symlinks` for the interpreter. This is separate from the
 toolchain_aliases repo because referencing the `python` interpreter target from
 this repo causes an eager fetch of the toolchain for the host platform.
-    """,
+
+This repo has three different ways in which is it called:
+
+1. Workspace. The `platforms` attribute is set, which are keys into the
+   PLATFORMS global. It assumes `name` + <matching platform name> is a
+   valid repo name which it can use as the backing repo.
+
+2. Bzlmod, created along side when python_register_toolchains is called
+   and expected to use one of repos created as part of that
+   python_register_toolchains call.
+   Because the bzlmod extension decides the platform mapping, it is given
+   the `platform`, `os_names`, and `archs` attributes to figure out which
+   to use.
+
+3. Bzlmod, created when the initial python_register_toolchains didn't
+   have a host-compatible runtime, so a different host-compatible
+   implementation was used.
+   This is like the normal bzlmod creation, except the python_versions
+   and suffixes may vary between choices, so the `impl_repo_suffixes` and
+   `python_version` attributes are specified.
+""",
     attrs = {
         "arch_names": attr.string_dict(
             doc = """
@@ -398,8 +418,43 @@ If set, overrides the platform metadata. Keyed by index in `platforms`
 If set, overrides the platform metadata. Keyed by index in `platforms`
 """,
         ),
-        "platforms": attr.string_list(mandatory = True),
-        "python_version": attr.string(mandatory = True),
+        "platforms": attr.string_list(
+            mandatory = True,
+            doc = """
+Platform names and backing repo-suffix.
+
+NOTE: The order of this list matters. The first platform that is compatible
+with the host will be selected; this can be customized by using the
+`RULES_PYTHON_REPO_TOOLCHAIN_*` env vars.
+
+When os_names aren't set, they act the key into the PLATFORMS
+dict to determine if a platform is compatible with the host. When
+os_names is set, then it is a (mostly) arbitrary platform name string
+(and platform metadata comes from the os_names/archs args).
+
+The string is used as a suffix to create the name of the repo that
+should be pointed to. i.e. `name` + <selected platform string> should
+result in a valid repo (e.g. created by python_register_toolchains()).
+Under bzlmod, this also means the same extension must create the
+repo named `name+suffix` and the host_toolchain repo.
+""",
+        ),
+        "python_version": attr.string(
+            mandatory = True,
+            doc = "Full python version, Major.Minor.Micro",
+        ),
+        "python_versions": attr.string_dict(
+            doc = """
+If set, the Python version for the corresponding selected platform.
+Keyed by index in `platforms`. Values Major.Minor.Patch
+""",
+        ),
+        "impl_repo_suffixes": attr.string_dict(
+            doc = """
+If set, the suffix to append to `name` to identify the backing repo that is used.
+Keyed by index in `platforms`.
+""",
+        ),
         "_rule_name": attr.string(default = "host_toolchain"),
         "_rules_python_workspace": attr.label(default = Label("//:WORKSPACE")),
     },
@@ -435,8 +490,8 @@ multi_toolchain_aliases = repository_rule(
     },
 )
 
-def sorted_host_platforms(platform_map):
-    """Sort the keys in the platform map to give correct precedence.
+def sorted_host_platform_names(platform_names):
+    """Sort platform names to give correct precedence.
 
     The order of keys in the platform mapping matters for the host toolchain
     selection. When multiple runtimes are compatible with the host, we take the
@@ -452,6 +507,29 @@ def sorted_host_platforms(platform_map):
     in a dict that autoformatters like to clobber and whose only documentation
     is an innocous looking formatter disable directive.
 
+    Args:
+        platform_names: a list of platform names
+
+    Returns:
+        list[str] the same values, but in the desired order.
+    """
+
+    def platform_keyer(name):
+        # Ascending sort: lower is higher precedence
+        pref = 0
+        if name.endswith("-" + FREETHREADED):
+            pref = 1
+        elif name.endswith("-" + MUSL):
+            pref = 2
+        return (pref, name)
+
+    return sorted(platform_map.keys(), key = platform_keyer)
+
+def sorted_host_platforms(platform_map):
+    """Sort the keys in the platform map to give correct precedence.
+
+    See sorted_host_platform_names for explanation.
+
     Args:
         platform_map: a mapping of platforms and their metadata.
 
@@ -467,13 +545,12 @@ def sorted_host_platforms(platform_map):
             1 if FREETHREADED in name else 0,
         )
 
-    sorted_platform_keys = sorted(platform_map.keys(), key = platform_keyer)
     return {
         key: platform_map[key]
-        for key in sorted_platform_keys
+        for key in sorted_host_platform_names(platform_map.keys())
     }
 
-def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platforms):
+def _get_host_impl_repo_suffix(*, rctx, logger, python_version, os_name, cpu_name, platforms):
     """Gets the host platform.
 
     Args:
@@ -493,6 +570,8 @@ def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platf
             platform_map[platform_name] = struct(
                 os_name = rctx.attr.os_names[key],
                 arch = rctx.attr.arch_names[key],
+                python_version = rctx.attr.python_versions.get(key),
+                impl_repo_suffix = rctx.attr.impl_repo_suffixes.get(key),
             )
     else:
         platform_map = sorted_host_platforms(PLATFORMS)
@@ -502,11 +581,13 @@ def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platf
         meta = platform_map[platform]
 
         if meta.os_name == os_name and meta.arch == cpu_name:
-            candidates.append(platform)
+            candidates.append((platform, meta))
 
     if len(candidates) == 1:
-        return candidates[0]
+        platform_name, meta = candidates[0]
+        return getattr(meta, "impl_repo_suffix", platform_name)
 
+    # todo: have this handle multiple python versions
     if candidates:
         env_var = "RULES_PYTHON_REPO_TOOLCHAIN_{}_{}_{}".format(
             python_version.replace(".", "_"),
@@ -520,12 +601,21 @@ def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platf
                 candidates,
             ))
         elif preference not in candidates:
+            # todo: need to map names like 3_13_0_linux_x86_64 back to
+            # the input values. Ah, er, wait
+            # Is this working?
+            # The return value is appended to this repo's name.
+            # This repo's name is e.g. python_3_13.
+            # the net result would be
+            # python_3_10_3_13_0_linux_x86_64
+            # which isn't a valid name
             return logger.fail("Please choose a preferred interpreter out of the following platforms: {}".format(candidates))
         else:
             candidates = [preference]
 
     if candidates:
-        return candidates[0]
+        platform_name, meta = candidates[0]
+        return getattr(meta, "impl_repo_suffix", platform_name)
 
     return logger.fail("Could not find a compatible 'host' python for '{os_name}', '{cpu_name}' from the loaded platforms: {platforms}".format(
         os_name = os_name,