feat(bzlmod): Allowing multiple python.toolchain extension calls

We do this work for two reasons. First, we must support Module
dependencies and sub-modules using python.toolchain. Second, we
needed this commit to supporting using multiple toolchains with bzlmod.

This commit modifies the python.toolchain extension to handle being
called multiple times. We are modeling how the multiple Python toolchains
work.

This commit implements various business logic in the toolchain class as
follows.

Toolchains in Sub Modules

It will create a toolchain in a sub-module if the toolchain
of the same name does not exist in the root module. The extension stops name
clashing between toolchains in the root module and sub-modules.
You cannot configure more than one toolchain as the default toolchain.

Toolchain set as the default version.

This extension will not create a toolchain in a sub-module
if the sub-module toolchain is marked as the default version. If you have
more than one toolchain in your root module, you need to set one of the
toolchains as the default version. If there is only one toolchain, it
is set as the default toolchain.

examples/bzlmod/MODULE.bazel

@@ -11,17 +11,21 @@ local_path_override(
     path = "../..",
 )
 
+PYTHON_NAME = "python3_9"
+
+PYTHON_TOOLCHAINS = PYTHON_NAME + "_toolchains"
+
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
-    name = "python3_9",
+    name = PYTHON_NAME,
     configure_coverage_tool = True,
     python_version = "3.9",
 )
-use_repo(python, "python3_9")
-use_repo(python, "python3_9_toolchains")
+use_repo(python, PYTHON_NAME)
+use_repo(python, PYTHON_TOOLCHAINS)
 
 register_toolchains(
-    "@python3_9_toolchains//:all",
+    "@{}//:all".format(PYTHON_TOOLCHAINS),
 )
 
 interpreter = use_extension("@rules_python//python/extensions:interpreter.bzl", "interpreter")
@@ -34,7 +38,8 @@ use_repo(interpreter, "interpreter_python3_9")
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 pip.parse(
     name = "pip",
-    # Intentionally set it false because the "true" case is already covered by examples/bzlmod_build_file_generation
+    # Intentionally set it false because the "true" case is already
+    # covered by examples/bzlmod_build_file_generation
     incompatible_generate_aliases = False,
     python_interpreter_target = "@interpreter_python3_9//:python",
     requirements_lock = "//:requirements_lock.txt",

examples/bzlmod_build_file_generation/MODULE.bazel

@@ -48,6 +48,8 @@ python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 # We also use the same name for python.host_python_interpreter.
 PYTHON_NAME = "python3"
 
+PYTHON_TOOLCHAINS = PYTHON_NAME + "_toolchains"
+
 # We next initialize the python toolchain using the extension.
 # You can set different Python versions in this block.
 python.toolchain(
@@ -63,12 +65,12 @@ python.toolchain(
 # into the scope of the current module.
 # All of the python3 repositories use the PYTHON_NAME as there prefix.  They
 # are not catenated for ease of reading.
-use_repo(python, PYTHON_NAME, "python3_toolchains")
+use_repo(python, PYTHON_NAME, PYTHON_TOOLCHAINS)
 
 # Register an already-defined toolchain so that Bazel can use it during
 # toolchain resolution.
 register_toolchains(
-    "@python3_toolchains//:all",
+    "@{}//:all".format(PYTHON_TOOLCHAINS),
 )
 
 # The interpreter extension discovers the platform specific Python binary.

python/extensions/python.bzl

@@ -17,40 +17,132 @@
 load("@rules_python//python:repositories.bzl", "python_register_toolchains")
 load("@rules_python//python/extensions/private:interpreter_hub.bzl", "hub_repo")
 
+# Printing a warning msg not debugging, so we have to disable
+# the buildifier check.
+# buildifier: disable=print
+def _print_warn(msg):
+    print(msg)
+
+def _python_register_toolchains(toolchain_attr, version_constraint):
+    python_register_toolchains(
+        name = toolchain_attr.name,
+        python_version = toolchain_attr.python_version,
+        register_coverage_tool = toolchain_attr.configure_coverage_tool,
+        ignore_root_user_error = toolchain_attr.ignore_root_user_error,
+        set_python_version_constraint = version_constraint,
+    )
+
 def _python_impl(module_ctx):
-    toolchains = []
+    # We collect all of the toolchain names to create
+    # the INTERPRETER_LABELS map.  This is used
+    # by interpreter_extensions.bzl via the hub_repo call below.
+    toolchain_names = []
+
+    # Used to store the default toolchain so we can create it last.
+    default_toolchain = None
+
+    # Used to store toolchains that are in sub modules.
+    sub_toolchains_map = {}
+
     for mod in module_ctx.modules:
         for toolchain_attr in mod.tags.toolchain:
-            python_register_toolchains(
-                name = toolchain_attr.name,
-                python_version = toolchain_attr.python_version,
-                bzlmod = True,
-                # Toolchain registration in bzlmod is done in MODULE file
-                register_toolchains = False,
-                register_coverage_tool = toolchain_attr.configure_coverage_tool,
-                ignore_root_user_error = toolchain_attr.ignore_root_user_error,
-            )
-
-            # We collect all of the toolchain names to create
-            # the INTERPRETER_LABELS map.  This is used
-            # by interpreter_extensions.bzl
-            toolchains.append(toolchain_attr.name)
+            # If we are in the root module we always register the toolchain.
+            # We wait to register the default toolchain till the end.
+            if mod.is_root:
+                toolchain_names.append(toolchain_attr.name)
+
+                # If we have the default version or we only have one toolchain
+                # in the root module we set the toolchain as the default toolchain.
+                if toolchain_attr.default_version or len(mod.tags.toolchain) == 1:
+                    # We have already found one default toolchain, and we can
+                    # only have one.
+                    if default_toolchain != None:
+                        fail("""We found more than one toolchain that is marked 
+as the default version.  Only set one toolchain with default_version set as 
+True.""")
+                    default_toolchain = toolchain_attr
+                    continue
+
+                #  Always register toolchains that are in the root module.
+                _python_register_toolchains(toolchain_attr, True)
+            else:
+                # We add the toolchain to a map, and we later create the
+                # toolchain if the root module does not have a toolchain with
+                # the same name.  We have to loop through all of the modules to
+                # ensure that we get a full list of the root toolchains.
+                sub_toolchains_map[toolchain_attr.name] = toolchain_attr
+
+    # We did not find a default toolchain so we fail.
+    if default_toolchain == None:
+        fail("""Unable to find a default toolchain in the root module.  
+Please define a toolchain that has default_version set to True.""")
+
+    # Create the toolchains in the submodule(s).
+    for name, toolchain_attr in sub_toolchains_map:
+        # A sub module cannot have a toolchain that is marked as the
+        # default version. TODO: should we create the toolchain anyways,
+        # but set the default version to False?
+        if toolchain_attr.default_version:
+            fail("""Not able to create toolchain named: {}.  This toolchain exists
+in a sub module and defalult_version is set to True.""".format(name))
 
+        # We cannot have a toolchain in a sub module that has the same name of
+        # a toolchain in the root module. This will cause name clashing.
+        if name in toolchain_names:
+            _print_warn("""Not creating the toolchain from sub module, with the name {}. The root
+ modhas a toolchain of the same name.""".format(toolchain_attr.name))
+            continue
+        toolchain_names.append(name)
+        _python_register_toolchains(toolchain_attr, True)
+
+    # We register the default toolchain last.
+    _python_register_toolchains(default_toolchain, False)
+
+    # Create the hub for the different interpreter versions.
     hub_repo(
         name = "pythons_hub",
-        toolchains = toolchains,
+        toolchains = toolchain_names,
     )
 
 python = module_extension(
-    doc = "Bzlmod extension that is used to register a Python toolchain.",
+    doc = """Bzlmod extension that is used to register Python toolchains.
+""",
     implementation = _python_impl,
     tag_classes = {
         "toolchain": tag_class(
+            doc = """Tag class used to register Python toolchains.
+Use this tag class to register one of more Python toolchains. This class
+is also potentially called by sub modules. The following covers different
+business rules and use cases.
+
+Toolchains in the Root Module
+
+This class registers all toolchains in the root module.
+
+Toolchains in Sub Modules
+
+It will create a toolchain that is in a sub module, if the toolchain
+of the same name does not exist in the root module.  The extension stops name
+clashing between toolchains in the root module and toolchains in sub modules.
+You cannot configure more than one toolchain as the default toolchain.  
+
+Toolchain set as the default versions
+
+This extension will not create a toolchain that exists in a sub module, 
+if the sub module toolchain is marked as the default version. If you have
+more than one toolchain in your root module, you need to set one of the
+toolchains as the default version.  If there is only one toolchain it 
+is set as the default toolchain.  
+""",
             attrs = {
                 "configure_coverage_tool": attr.bool(
                     mandatory = False,
                     doc = "Whether or not to configure the default coverage tool for the toolchains.",
                 ),
+                "default_version": attr.bool(
+                    mandatory = False,
+                    doc = "Whether the toolchain is the default version",
+                ),
                 "ignore_root_user_error": attr.bool(
                     default = False,
                     doc = "Whether the check for root should be ignored or not. This causes cache misses with .pyc files.",

python/repositories.bzl

@@ -463,7 +463,6 @@ def python_register_toolchains(
         register_coverage_tool = False,
         set_python_version_constraint = False,
         tool_versions = TOOL_VERSIONS,
-        bzlmod = False,
         **kwargs):
     """Convenience macro for users which does typical setup.
 
@@ -486,9 +485,15 @@ def python_register_toolchains(
         set_python_version_constraint: When set to true, target_compatible_with for the toolchains will include a version constraint.
         tool_versions: a dict containing a mapping of version with SHASUM and platform info. If not supplied, the defaults
             in python/versions.bzl will be used.
-        bzlmod: Whether this rule is being run under a bzlmod module extension.
         **kwargs: passed to each python_repositories call.
     """
+
+    # If we have @@ we have bzlmod
+    bzlmod = str(Label("//:distribution")).startswith("@@")
+    if bzlmod:
+        # you cannot used native.register_toolchains when using bzlmod.
+        register_toolchains = False
+
     base_url = kwargs.pop("base_url", DEFAULT_RELEASE_BASE_URL)
 
     if python_version in MINOR_MAPPING: