refactor: make bzlmod directly aware of created toolchain repo names

Author: rickeylev

python/private/python.bzl

@@ -267,11 +267,18 @@ def parse_modules(*, module_ctx, _fail = fail):
 def _python_impl(module_ctx):
     py = parse_modules(module_ctx = module_ctx)
 
-    # dict[str version, list[str] platforms]; where version is full
-    # python version string ("3.4.5"), and platforms are keys from
-    # the PLATFORMS global.
-    loaded_platforms = {}
-    for toolchain_info in py.toolchains:
+    # list of structs; see inline struct call within the loop below.
+    toolchain_impls = []
+
+    # list[str] of the base names of toolchain repos
+    base_toolchain_repo_names = []
+
+    # Create the underlying python_repository repos that contain the
+    # python runtimes and their toolchain implementation definitions.
+    for i, toolchain_info in enumerate(py.toolchains):
+        is_last = (i + 1) == len(py.toolchains)
+        base_toolchain_repo_names.append(toolchain_info.name)
+
         # Ensure that we pass the full version here.
         full_python_version = full_version(
             version = toolchain_info.python_version,
@@ -286,12 +293,28 @@ def _python_impl(module_ctx):
         kwargs.update(py.config.kwargs.get(toolchain_info.python_version, {}))
         kwargs.update(py.config.kwargs.get(full_python_version, {}))
         kwargs.update(py.config.default)
-        toolchain_registered_platforms = python_register_toolchains(
+        register_result = python_register_toolchains(
             name = toolchain_info.name,
             _internal_bzlmod_toolchain_call = True,
             **kwargs
         )
-        loaded_platforms[full_python_version] = toolchain_registered_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
+                name = repo_name,
+                # str: The repo name the toolchain() target points to.
+                impl_repo_name = repo_name,
+                # str: platform key in the passed-in platforms dict
+                platform_name = platform_name,
+                # struct: platform_info() struct
+                platform = platform_info,
+                # str: Major.Minor.Micro python version
+                full_python_version = full_python_version,
+                # bool: whether to implicitly add the python version constraint
+                # to the toolchain's target_settings.
+                # The last toolchain is the default; it can't have version constraints
+                set_python_version_constraint = is_last,
+            ))
 
     # List of the base names ("python_3_10") for the toolchain repos
     base_toolchain_repo_names = []
@@ -329,31 +352,23 @@ def _python_impl(module_ctx):
 
     # Split the toolchain info into separate objects so they can be passed onto
     # the repository rule.
-    for i, t in enumerate(py.toolchains):
-        is_last = (i + 1) == len(py.toolchains)
-        base_name = t.name
-        base_toolchain_repo_names.append(base_name)
-        fv = full_version(version = t.python_version, minor_mapping = py.config.minor_mapping)
-        platforms = loaded_platforms[fv]
-        for platform_name, platform_info in platforms.items():
-            key = str(len(toolchain_names))
-
-            full_name = "{}_{}".format(base_name, platform_name)
-            toolchain_names.append(full_name)
-            toolchain_repo_names[key] = full_name
-            toolchain_tcw_map[key] = platform_info.compatible_with
-
-            # The target_settings attribute may not be present for users
-            # patching python/versions.bzl.
-            toolchain_ts_map[key] = getattr(platform_info, "target_settings", [])
-            toolchain_platform_keys[key] = platform_name
-            toolchain_python_versions[key] = fv
-
-            # The last toolchain is the default; it can't have version constraints
-            # Despite the implication of the arg name, the values are strs, not bools
-            toolchain_set_python_version_constraints[key] = (
-                "True" if not is_last else "False"
-            )
+    for entry in toolchain_impls:
+        key = str(len(toolchain_names))
+
+        toolchain_names.append(entry.name)
+        toolchain_repo_names[key] = entry.impl_repo_name
+        toolchain_tcw_map[key] = entry.platform.compatible_with
+
+        # The target_settings attribute may not be present for users
+        # patching python/versions.bzl.
+        toolchain_ts_map[key] = getattr(entry.platform, "target_settings", [])
+        toolchain_platform_keys[key] = entry.platform_name
+        toolchain_python_versions[key] = entry.full_python_version
+
+        # Repo rules can't accept dict[str, bool], so encode them as a string value.
+        toolchain_set_python_version_constraints[key] = (
+            "True" if entry.set_python_version_constraint else "False"
+        )
 
     hub_repo(
         name = "pythons_hub",

python/private/python_register_toolchains.bzl

@@ -111,13 +111,17 @@ def python_register_toolchains(
                 ))
             register_coverage_tool = False
 
-    loaded_platforms = {}
-    for platform in platforms.keys():
+    # list[str] of the platform names that were used
+    loaded_platforms = []
+
+    # dict[str repo name, tuple[str, platform_info]]
+    impl_repos = {}
+    for platform, platform_info in platforms.items():
         sha256 = tool_versions[python_version]["sha256"].get(platform, None)
         if not sha256:
             continue
 
-        loaded_platforms[platform] = platforms[platform]
+        loaded_platforms.append(platform)
         (release_filename, urls, strip_prefix, patches, patch_strip) = get_release_info(platform, python_version, base_url, tool_versions)
 
         # allow passing in a tool version
@@ -137,11 +141,10 @@ def python_register_toolchains(
                 )],
             )
 
+        impl_repo_name = "{}_{}".format(name, platform)
+        impl_repos[impl_repo_name] = (platform, platform_info)
         python_repository(
-            name = "{name}_{platform}".format(
-                name = name,
-                platform = platform,
-            ),
+            name = impl_repo_name,
             sha256 = sha256,
             patches = patches,
             patch_strip = patch_strip,
@@ -169,26 +172,29 @@ def python_register_toolchains(
 
     host_toolchain(
         name = name + "_host",
-        platforms = loaded_platforms.keys(),
+        platforms = loaded_platforms,
         python_version = python_version,
     )
 
     toolchain_aliases(
         name = name,
         python_version = python_version,
         user_repository_name = name,
-        platforms = loaded_platforms.keys(),
+        platforms = loaded_platforms,
     )
 
     # in bzlmod we write out our own toolchain repos
     if bzlmod_toolchain_call:
-        return loaded_platforms
+        return struct(
+            # dict[str name, tuple[str platform_name, platform_info]]
+            impl_repos = impl_repos,
+        )
 
     toolchains_repo(
         name = toolchain_repo_name,
         python_version = python_version,
         set_python_version_constraint = set_python_version_constraint,
         user_repository_name = name,
-        platforms = loaded_platforms.keys(),
+        platforms = loaded_platforms,
     )
     return None