refactor: explicitly define host platform ordering selection

Author: rickeylev

python/private/python.bzl

@@ -21,7 +21,7 @@ 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")
+load(":toolchains_repo.bzl", "host_toolchain", "multi_toolchain_aliases", "sorted_host_platforms")
 load(":util.bzl", "IS_BAZEL_6_4_OR_HIGHER")
 load(":version.bzl", "version")
 
@@ -298,9 +298,8 @@ def _python_impl(module_ctx):
             _internal_bzlmod_toolchain_call = True,
             **kwargs
         )
-        host_platforms = []
-        host_os_names = {}
-        host_archs = {}
+
+        host_platforms = {}
         for repo_name, (platform_name, platform_info) in register_result.impl_repos.items():
             toolchain_impls.append(struct(
                 # str: The base name to use for the toolchain() target
@@ -319,17 +318,21 @@ def _python_impl(module_ctx):
                 set_python_version_constraint = is_last,
             ))
             if _is_compatible_with_host(module_ctx, platform_info):
-                host_key = str(len(host_platforms))
-                host_platforms.append(platform_name)
-                host_os_names[host_key] = platform_info.os_name
-                host_archs[host_key] = platform_info.arch
+                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,
-            os_names = host_os_names,
-            archs = host_archs,
+            platforms = host_platforms.keys(),
+            os_names = {
+                str(i): platform_info.os_name
+                for i, platform_info in enumerate(host_platforms.values())
+            },
+            archs = {
+                str(i): platform_info.arch
+                for i, platform_info in enumerate(host_platforms.values())
+            },
             python_version = full_python_version,
         )
 

python/private/toolchains_repo.bzl

@@ -25,6 +25,8 @@ platform-specific repositories.
 
 load(
     "//python:versions.bzl",
+    "FREETHREADED",
+    "MUSL",
     "PLATFORMS",
     "WINDOWS_NAME",
 )
@@ -386,18 +388,18 @@ toolchain_aliases repo because referencing the `python` interpreter target from
 this repo causes an eager fetch of the toolchain for the host platform.
     """,
     attrs = {
-        "platforms": attr.string_list(mandatory = True),
-        "python_version": attr.string(mandatory = True),
-        "os_names": attr.string_dict(
+        "archs": attr.string_dict(
             doc = """
 If set, overrides the platform metadata. Keyed by index in `platforms`
 """,
         ),
-        "archs": attr.string_dict(
+        "os_names": attr.string_dict(
             doc = """
 If set, overrides the platform metadata. Keyed by index in `platforms`
 """,
         ),
+        "platforms": attr.string_list(mandatory = True),
+        "python_version": attr.string(mandatory = True),
         "_rule_name": attr.string(default = "host_toolchain"),
         "_rules_python_workspace": attr.label(default = Label("//:WORKSPACE")),
     },
@@ -436,6 +438,46 @@ multi_toolchain_aliases = repository_rule(
 def sanitize_platform_name(platform):
     return platform.replace("-", "_")
 
+def sorted_host_platforms(platform_map):
+    """Sort the keys in the platform map 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
+    first that is compatible (usually; there's also the
+    `RULES_PYTHON_REPO_TOOLCHAIN_*` environment variables). The historical
+    behavior carefully constructed the ordering of platform keys such that
+    the ordering was:
+    * Regular platforms
+    * The "-freethreaded" suffix
+    * The "-musl" suffix
+
+    Here, we formalize that so it isn't subtly encoded in the ordering of keys
+    in a dict that autoformatters like to clobber and whose only documentation
+    is an innocous looking formatter disable directive.
+
+    Args:
+        platform_map: a mapping of platforms and their metadata.
+
+    Returns:
+        dict; the same values, but with the keys inserted in the desired
+        order so that iteration happens 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)
+
+    sorted_platform_keys = sorted(platform_map.keys(), key = platform_keyer)
+    return {
+        key: platform_map[key]
+        for key in sorted_platform_keys
+    }
+
 def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platforms):
     """Gets the host platform.
 
@@ -458,7 +500,7 @@ def _get_host_platform(*, rctx, logger, python_version, os_name, cpu_name, platf
                 arch = rctx.attr.archs[key],
             )
     else:
-        platform_map = PLATFORMS
+        platform_map = sorted_host_platforms(PLATFORMS)
 
     candidates = []
     for platform in platforms:

python/versions.bzl

@@ -19,7 +19,9 @@
 MACOS_NAME = "osx"
 LINUX_NAME = "linux"
 WINDOWS_NAME = "windows"
+
 FREETHREADED = "freethreaded"
+MUSL = "musl"
 INSTALL_ONLY = "install_only"
 
 DEFAULT_RELEASE_BASE_URL = "https://github.com/astral-sh/python-build-standalone/releases/download"