feat(bzlmod): Moving register.toolchains internal (#1238)

This commit moves the register.toolchains bzlmod call to inside
of rules_python.  Instead of a user having to call register.toolchains
in their MODULE.bazel, rules_python/MODULE.bazel calls it
on the internal hub.

This is a breaking change if you are using register.toolchains inside
of submodules.  Using register.toolchains inside of submodules is
not recommended anyways.  This is now broken because we are not
creating a repo for every Python version toolchain.  All of the
toochain calls exist now in the hub's repo BUILD.bazel file.

20 files changed, 579 insertions, 229 deletions

diff --git a/.bazelrc b/.bazelrc
index fe542b3..3c31741 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -3,8 +3,8 @@
 # This lets us glob() up all the files inside the examples to make them inputs to tests
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, run tools/bazel_integration_test/update_deleted_packages.sh
-build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_point,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
 
 test --test_output=errors
 
diff --git a/BZLMOD_SUPPORT.md b/BZLMOD_SUPPORT.md
index cf95d12..8efd0df 100644
--- a/BZLMOD_SUPPORT.md
+++ b/BZLMOD_SUPPORT.md
@@ -31,7 +31,7 @@ A second example, in [examples/bzlmod_build_file_generation](examples/bzlmod_bui
 
 This rule set does not have full feature partity with the older `WORKSPACE` type configuration:
 
-1. Multiple python versions are not yet supported, as demonstrated in [this](examples/multi_python_versions) example.
+1. Multiple pip extensions are not yet supported, as demonstrated in [this](examples/multi_python_versions) example.
 2. Gazelle does not support finding deps in sub-modules.  For instance we can have a dep like ` "@our_other_module//other_module/pkg:lib",` in a `py_test` definition.
 
 Check ["issues"](/bazelbuild/rules_python/issues) for an up to date list.
diff --git a/MODULE.bazel b/MODULE.bazel
index ddd946c..b45c2ff 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -47,5 +47,10 @@ use_repo(
     "pypi__coverage_cp39_x86_64-unknown-linux-gnu",
 )
 
+# We need to do another use_extension call to expose the "pythons_hub"
+# repo.
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 use_repo(python, "pythons_hub")
+
+# This call registers the Python toolchains.
+register_toolchains("@pythons_hub//:all")
diff --git a/README.md b/README.md
index a3f1886..6893a1d 100644
--- a/README.md
+++ b/README.md
@@ -53,32 +53,39 @@ To import rules_python in your project, you first need to add it to your
 
 To register a hermetic Python toolchain rather than rely on a system-installed interpreter for runtime execution, you can add to the `MODULE.bazel` file:
 
-```python
+```starlark
 # Find the latest version number here: https://github.com/bazelbuild/rules_python/releases
 # and change the version number if needed in the line below.
-bazel_dep(name = "rules_python", version = "0.20.0")
+bazel_dep(name = "rules_python", version = "0.21.0")
+
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(
+    name = "python",
+    configure_coverage_tool = True,
+    is_default = True,
+    python_version = "3.9",
+)
 
-# You do not have to use pip for the toolchain, but most people
-# will use it for the dependency management.
-pip = use_extension("@rules_python//python:extensions.bzl", "pip")
+interpreter = use_extension("@rules_python//python/extensions:interpreter.bzl", "interpreter")
+interpreter.install(
+    name = "interpreter",
+    python_name = "python",
+)
+use_repo(interpreter, "interpreter")
 
+pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 pip.parse(
     name = "pip",
+    incompatible_generate_aliases = True,
+    python_interpreter_target = "@interpreter//:python",
     requirements_lock = "//:requirements_lock.txt",
+    requirements_windows = "//:requirements_windows.txt",
 )
-
 use_repo(pip, "pip")
-
-# Register a specific python toolchain instead of using the host version
-python = use_extension("@rules_python//python:extensions.bzl", "python")
-
-use_repo(python, "python3_10_toolchains")
-
-register_toolchains(
-    "@python3_10_toolchains//:all",
-)
 ```
 
+For more documentation see the bzlmod examples under the [examples](examples) folder.
+
 ### Using a WORKSPACE file
 
 To import rules_python in your project, you first need to add it to your
diff --git a/examples/bzlmod/BUILD.bazel b/examples/bzlmod/BUILD.bazel
index 8649822..0a068ce 100644
--- a/examples/bzlmod/BUILD.bazel
+++ b/examples/bzlmod/BUILD.bazel
@@ -8,9 +8,6 @@
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load("@pip//:requirements.bzl", "all_requirements", "all_whl_requirements", "requirement")
 load("@python_39//:defs.bzl", py_test_with_transition = "py_test")
-
-# This is not working yet till the toolchain hub registration is working
-# load("@python_310//:defs.bzl", py_binary_310 = "py_binary")
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@rules_python//python:pip.bzl", "compile_pip_requirements")
 
@@ -50,22 +47,6 @@ py_binary(
     ],
 )
 
-# This is still WIP.  Not working till we have the toolchain
-# registration functioning.
-
-# This is used for testing mulitple versions of Python. This is
-# used only when you need to support multiple versions of Python
-# in the same project.
-# py_binary_310(
-#     name = "main_310",
-#     srcs = ["__main__.py"],
-#     main = "__main__.py",
-#     visibility = ["//:__subpackages__"],
-#     deps = [
-#         ":lib",
-#     ],
-# )
-
 # see https://bazel.build/reference/be/python#py_test
 py_test(
     name = "test",
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
index bb4183b..24bb458 100644
--- a/examples/bzlmod/MODULE.bazel
+++ b/examples/bzlmod/MODULE.bazel
@@ -15,14 +15,10 @@ local_path_override(
 # We also use the same value in the python.host_python_interpreter call.
 PYTHON_NAME_39 = "python_39"
 
-PYTHON_39_TOOLCHAINS = PYTHON_NAME_39 + "_toolchains"
-
 INTERPRETER_NAME_39 = "interpreter_39"
 
 PYTHON_NAME_310 = "python_310"
 
-PYTHON_310_TOOLCHAINS = PYTHON_NAME_310 + "_toolchains"
-
 INTERPRETER_NAME_310 = "interpreter_310"
 
 # We next initialize the python toolchain using the extension.
@@ -50,25 +46,9 @@ python.toolchain(
     python_version = "3.10",
 )
 
-# use_repo imports one or more repos generated by the given module extension
-# into the scope of the current module. We are importing the various repos
-# created by the above python.toolchain calls.
-use_repo(
-    python,
-    PYTHON_NAME_39,
-    PYTHON_39_TOOLCHAINS,
-    PYTHON_NAME_310,
-    PYTHON_310_TOOLCHAINS,
-)
-
-# This call registers the Python toolchains.
-# Note: there is work under way to move this code to within
-# rules_python, and the user won't have to make this call,
-# unless they are registering custom toolchains.
-register_toolchains(
-    "@{}//:all".format(PYTHON_39_TOOLCHAINS),
-    "@{}//:all".format(PYTHON_310_TOOLCHAINS),
-)
+# You only need to load this repositories if you are using muiltple Python versions.
+# See the tests folder for various examples.
+use_repo(python, PYTHON_NAME_39, "python_aliases")
 
 # The interpreter extension discovers the platform specific Python binary.
 # It creates a symlink to the binary, and we pass the label to the following
diff --git a/examples/bzlmod/other_module/MODULE.bazel b/examples/bzlmod/other_module/MODULE.bazel
index eebfbca..5fb7452 100644
--- a/examples/bzlmod/other_module/MODULE.bazel
+++ b/examples/bzlmod/other_module/MODULE.bazel
@@ -12,12 +12,8 @@ bazel_dep(name = "rules_python", version = "")
 # testing purposes.
 PYTHON_NAME_39 = "python_39"
 
-PYTHON_39_TOOLCHAINS = PYTHON_NAME_39 + "_toolchains"
-
 PYTHON_NAME_311 = "python_311"
 
-PYTHON_311_TOOLCHAINS = PYTHON_NAME_311 + "_toolchains"
-
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
     # This name is used in the various use_repo statements
@@ -42,16 +38,5 @@ python.toolchain(
 use_repo(
     python,
     PYTHON_NAME_39,
-    PYTHON_39_TOOLCHAINS,
     PYTHON_NAME_311,
-    PYTHON_311_TOOLCHAINS,
-)
-
-# This call registers the Python toolchains.
-# Note: there is work under way to move this code to within
-# rules_python, and the user won't have to make this call,
-# unless they are registering custom toolchains.
-register_toolchains(
-    "@{}//:all".format(PYTHON_39_TOOLCHAINS),
-    "@{}//:all".format(PYTHON_311_TOOLCHAINS),
 )
diff --git a/examples/bzlmod/tests/BUILD.bazel b/examples/bzlmod/tests/BUILD.bazel
new file mode 100644
index 0000000..5331f4a
--- /dev/null
+++ b/examples/bzlmod/tests/BUILD.bazel
@@ -0,0 +1,142 @@
+load("@python_aliases//3.10:defs.bzl", py_binary_3_10 = "py_binary", py_test_3_10 = "py_test")
+load("@python_aliases//3.11:defs.bzl", py_binary_3_11 = "py_binary", py_test_3_11 = "py_test")
+load("@python_aliases//3.9:defs.bzl", py_binary_3_9 = "py_binary", py_test_3_9 = "py_test")
+load("@rules_python//python:defs.bzl", "py_binary", "py_test")
+
+py_binary(
+    name = "version_default",
+    srcs = ["version.py"],
+    main = "version.py",
+)
+
+py_binary_3_9(
+    name = "version_3_9",
+    srcs = ["version.py"],
+    main = "version.py",
+)
+
+py_binary_3_10(
+    name = "version_3_10",
+    srcs = ["version.py"],
+    main = "version.py",
+)
+
+py_binary_3_11(
+    name = "version_3_11",
+    srcs = ["version.py"],
+    main = "version.py",
+)
+
+# This is a work in progress and the commented
+# tests will not work  until we can support
+# multiple pips with bzlmod.
+
+#py_test(
+#    name = "my_lib_default_test",
+#    srcs = ["my_lib_test.py"],
+#    main = "my_lib_test.py",
+#    deps = ["//libs/my_lib"],
+#)
+
+#py_test_3_9(
+#    name = "my_lib_3_9_test",
+#    srcs = ["my_lib_test.py"],
+#    main = "my_lib_test.py",
+#    deps = ["//libs/my_lib"],
+#)
+
+#py_test_3_10(
+#    name = "my_lib_3_10_test",
+#    srcs = ["my_lib_test.py"],
+#    main = "my_lib_test.py",
+#    deps = ["//libs/my_lib"],
+#)
+
+#py_test_3_11(
+#    name = "my_lib_3_11_test",
+#    srcs = ["my_lib_test.py"],
+#    main = "my_lib_test.py",
+#    deps = ["//libs/my_lib"],
+#)
+
+py_test(
+    name = "version_default_test",
+    srcs = ["version_test.py"],
+    env = {"VERSION_CHECK": "3.9"},  # The default defined in the WORKSPACE.
+    main = "version_test.py",
+)
+
+py_test_3_9(
+    name = "version_3_9_test",
+    srcs = ["version_test.py"],
+    env = {"VERSION_CHECK": "3.9"},
+    main = "version_test.py",
+)
+
+py_test_3_10(
+    name = "version_3_10_test",
+    srcs = ["version_test.py"],
+    env = {"VERSION_CHECK": "3.10"},
+    main = "version_test.py",
+)
+
+py_test_3_11(
+    name = "version_3_11_test",
+    srcs = ["version_test.py"],
+    env = {"VERSION_CHECK": "3.11"},
+    main = "version_test.py",
+)
+
+py_test(
+    name = "version_default_takes_3_10_subprocess_test",
+    srcs = ["cross_version_test.py"],
+    data = [":version_3_10"],
+    env = {
+        "SUBPROCESS_VERSION_CHECK": "3.10",
+        "SUBPROCESS_VERSION_PY_BINARY": "$(rootpath :version_3_10)",
+        "VERSION_CHECK": "3.9",
+    },
+    main = "cross_version_test.py",
+)
+
+py_test_3_10(
+    name = "version_3_10_takes_3_9_subprocess_test",
+    srcs = ["cross_version_test.py"],
+    data = [":version_3_9"],
+    env = {
+        "SUBPROCESS_VERSION_CHECK": "3.9",
+        "SUBPROCESS_VERSION_PY_BINARY": "$(rootpath :version_3_9)",
+        "VERSION_CHECK": "3.10",
+    },
+    main = "cross_version_test.py",
+)
+
+sh_test(
+    name = "version_test_binary_default",
+    srcs = ["version_test.sh"],
+    data = [":version_default"],
+    env = {
+        "VERSION_CHECK": "3.9",  # The default defined in the WORKSPACE.
+        "VERSION_PY_BINARY": "$(rootpath :version_default)",
+    },
+)
+
+sh_test(
+    name = "version_test_binary_3_9",
+    srcs = ["version_test.sh"],
+    data = [":version_3_9"],
+    env = {
+        "VERSION_CHECK": "3.9",
+        "VERSION_PY_BINARY": "$(rootpath :version_3_9)",
+    },
+)
+
+sh_test(
+    name = "version_test_binary_3_10",
+    srcs = ["version_test.sh"],
+    data = [":version_3_10"],
+    env = {
+        "VERSION_CHECK": "3.10",
+        "VERSION_PY_BINARY": "$(rootpath :version_3_10)",
+    },
+)
diff --git a/examples/bzlmod/tests/cross_version_test.py b/examples/bzlmod/tests/cross_version_test.py
new file mode 100644
index 0000000..437be2e
--- /dev/null
+++ b/examples/bzlmod/tests/cross_version_test.py
@@ -0,0 +1,39 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import sys
+
+process = subprocess.run(
+    [os.getenv("SUBPROCESS_VERSION_PY_BINARY")],
+    stdout=subprocess.PIPE,
+    universal_newlines=True,
+)
+
+subprocess_current = process.stdout.strip()
+subprocess_expected = os.getenv("SUBPROCESS_VERSION_CHECK")
+
+if subprocess_current != subprocess_expected:
+    print(
+        f"expected subprocess version '{subprocess_expected}' is different than returned '{subprocess_current}'"
+    )
+    sys.exit(1)
+
+expected = os.getenv("VERSION_CHECK")
+current = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+if current != expected:
+    print(f"expected version '{expected}' is different than returned '{current}'")
+    sys.exit(1)
diff --git a/examples/bzlmod/tests/version.py b/examples/bzlmod/tests/version.py
new file mode 100644
index 0000000..2d293c1
--- /dev/null
+++ b/examples/bzlmod/tests/version.py
@@ -0,0 +1,17 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+print(f"{sys.version_info.major}.{sys.version_info.minor}")
diff --git a/examples/bzlmod/tests/version_test.py b/examples/bzlmod/tests/version_test.py
new file mode 100644
index 0000000..444f5e4
--- /dev/null
+++ b/examples/bzlmod/tests/version_test.py
@@ -0,0 +1,23 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+
+expected = os.getenv("VERSION_CHECK")
+current = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+if current != expected:
+    print(f"expected version '{expected}' is different than returned '{current}'")
+    sys.exit(1)
diff --git a/examples/bzlmod/tests/version_test.sh b/examples/bzlmod/tests/version_test.sh
new file mode 100755
index 0000000..3bedb95
--- /dev/null
+++ b/examples/bzlmod/tests/version_test.sh
@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+set -o errexit -o nounset -o pipefail
+
+version_py_binary=$("${VERSION_PY_BINARY}")
+
+if [[ "${version_py_binary}" != "${VERSION_CHECK}" ]]; then
+    echo >&2 "expected version '${VERSION_CHECK}' is different than returned '${version_py_binary}'"
+    exit 1
+fi
diff --git a/examples/bzlmod_build_file_generation/BUILD.bazel b/examples/bzlmod_build_file_generation/BUILD.bazel
index 05a15cc..498969b 100644
--- a/examples/bzlmod_build_file_generation/BUILD.bazel
+++ b/examples/bzlmod_build_file_generation/BUILD.bazel
@@ -7,7 +7,6 @@
 # requirements.
 load("@bazel_gazelle//:def.bzl", "gazelle")
 load("@pip//:requirements.bzl", "all_whl_requirements")
-load("@python//:defs.bzl", py_test_with_transition = "py_test")
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 load("@rules_python//python:pip.bzl", "compile_pip_requirements")
 load("@rules_python_gazelle_plugin//:def.bzl", "GAZELLE_PYTHON_RUNTIME_DEPS")
@@ -72,13 +71,6 @@ gazelle(
     gazelle = "@rules_python_gazelle_plugin//python:gazelle_binary",
 )
 
-py_test_with_transition(
-    name = "test_with_transition",
-    srcs = ["__test__.py"],
-    main = "__test__.py",
-    deps = [":bzlmod_build_file_generation"],
-)
-
 # The following targets are created and maintained by gazelle
 py_library(
     name = "bzlmod_build_file_generation",
diff --git a/examples/bzlmod_build_file_generation/MODULE.bazel b/examples/bzlmod_build_file_generation/MODULE.bazel
index 45a1318..d69dd7d 100644
--- a/examples/bzlmod_build_file_generation/MODULE.bazel
+++ b/examples/bzlmod_build_file_generation/MODULE.bazel
@@ -48,38 +48,17 @@ python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 # We also use the same name for python.host_python_interpreter.
 PYTHON_NAME = "python"
 
-PYTHON_TOOLCHAINS = PYTHON_NAME + "_toolchains"
-
 INTERPRETER_NAME = "interpreter"
 
 # We next initialize the python toolchain using the extension.
 # You can set different Python versions in this block.
 python.toolchain(
-    # This name is used in the various use_repo statements
-    # below, and in the local extension that is in this
-    # example.
     name = PYTHON_NAME,
     configure_coverage_tool = True,
     is_default = True,
     python_version = "3.9",
 )
 
-# Import the python repositories generated by the given module extension
-# 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,
-    PYTHON_TOOLCHAINS,
-)
-
-# Register an already-defined toolchain so that Bazel can use it during
-# toolchain resolution.
-register_toolchains(
-    "@{}//:all".format(PYTHON_TOOLCHAINS),
-)
-
 # The interpreter extension discovers the platform specific Python binary.
 # It creates a symlink to the binary, and we pass the label to the following
 # pip.parse call.
diff --git a/examples/py_proto_library/MODULE.bazel b/examples/py_proto_library/MODULE.bazel
index 6fb1a05..3116c40 100644
--- a/examples/py_proto_library/MODULE.bazel
+++ b/examples/py_proto_library/MODULE.bazel
@@ -18,11 +18,7 @@ python.toolchain(
     configure_coverage_tool = True,
     python_version = "3.9",
 )
-use_repo(python, "python3_9_toolchains")
-
-register_toolchains(
-    "@python3_9_toolchains//:all",
-)
+use_repo(python, "python3_9")
 
 # We are using rules_proto to define rules_proto targets to be consumed by py_proto_library.
 bazel_dep(name = "rules_proto", version = "5.3.0-21.7")
diff --git a/python/extensions/private/interpreter_hub.bzl b/python/extensions/private/interpreter_hub.bzl
deleted file mode 100644
index 82fcbf6..0000000
--- a/python/extensions/private/interpreter_hub.bzl
+++ /dev/null
@@ -1,69 +0,0 @@
-# Copyright 2023 The Bazel Authors. All rights reserved
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"Repo rule used by bzlmod extension to create a repo that has a map of Python interpreters and their labels"
-
-load("//python:versions.bzl", "WINDOWS_NAME")
-load("//python/private:toolchains_repo.bzl", "get_host_os_arch", "get_host_platform")
-
-_build_file_for_hub_template = """
-INTERPRETER_LABELS = {{
-{interpreter_labels}
-}}
-DEFAULT_TOOLCHAIN_NAME = "{default}"
-"""
-
-_line_for_hub_template = """\
-    "{name}": Label("@{name}_{platform}//:{path}"),
-"""
-
-def _hub_repo_impl(rctx):
-    (os, arch) = get_host_os_arch(rctx)
-    platform = get_host_platform(os, arch)
-
-    rctx.file("BUILD.bazel", "")
-    is_windows = (os == WINDOWS_NAME)
-    path = "python.exe" if is_windows else "bin/python3"
-
-    interpreter_labels = "\n".join([_line_for_hub_template.format(
-        name = name,
-        platform = platform,
-        path = path,
-    ) for name in rctx.attr.toolchains])
-
-    rctx.file(
-        "interpreters.bzl",
-        _build_file_for_hub_template.format(
-            interpreter_labels = interpreter_labels,
-            default = rctx.attr.default_toolchain,
-        ),
-    )
-
-hub_repo = repository_rule(
-    doc = """\
-This private rule create a repo with a BUILD file that contains a map of interpreter names
-and the labels to said interpreters. This map is used to by the interpreter hub extension.
-""",
-    implementation = _hub_repo_impl,
-    attrs = {
-        "default_toolchain": attr.string(
-            doc = "Name of the default toolchain",
-            mandatory = True,
-        ),
-        "toolchains": attr.string_list(
-            doc = "List of the base names the toolchain repo defines.",
-            mandatory = True,
-        ),
-    },
-)
diff --git a/python/extensions/private/pythons_hub.bzl b/python/extensions/private/pythons_hub.bzl
new file mode 100644
index 0000000..5baaef9
--- /dev/null
+++ b/python/extensions/private/pythons_hub.bzl
@@ -0,0 +1,136 @@
+# Copyright 2023 The Bazel Authors. All rights reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"Repo rule used by bzlmod extension to create a repo that has a map of Python interpreters and their labels"
+
+load("//python:versions.bzl", "MINOR_MAPPING", "WINDOWS_NAME")
+load(
+    "//python/private:toolchains_repo.bzl",
+    "get_host_os_arch",
+    "get_host_platform",
+    "get_repository_name",
+    "python_toolchain_build_file_content",
+)
+
+def _have_same_length(*lists):
+    if not lists:
+        fail("expected at least one list")
+    return len({len(length): None for length in lists}) == 1
+
+def _get_version(python_version):
+    # we need to get the MINOR_MAPPING or use the full version
+    if python_version in MINOR_MAPPING:
+        python_version = MINOR_MAPPING[python_version]
+    return python_version
+
+def _python_toolchain_build_file_content(
+        prefixes,
+        python_versions,
+        set_python_version_constraints,
+        user_repository_names,
+        workspace_location):
+    """This macro iterates over each of the lists and returns the toolchain content.
+
+    python_toolchain_build_file_content is called to generate each of the toolchain
+    definitions.
+    """
+
+    if not _have_same_length(python_versions, set_python_version_constraints, user_repository_names):
+        fail("all lists must have the same length")
+
+    rules_python = get_repository_name(workspace_location)
+
+    # Iterate over the length of python_versions and call
+    # build the toolchain content by calling python_toolchain_build_file_content
+    return "\n".join([python_toolchain_build_file_content(
+        prefix = prefixes[i],
+        python_version = _get_version(python_versions[i]),
+        set_python_version_constraint = set_python_version_constraints[i],
+        user_repository_name = user_repository_names[i],
+        rules_python = rules_python,
+    ) for i in range(len(python_versions))])
+
+_build_file_for_hub_template = """
+INTERPRETER_LABELS = {{
+{interpreter_labels}
+}}
+"""
+
+_line_for_hub_template = """\
+    "{name}": Label("@{name}_{platform}//:{path}"),
+"""
+
+def _hub_repo_impl(rctx):
+    # Create the various toolchain definitions and
+    # write them to the BUILD file.
+    rctx.file(
+        "BUILD.bazel",
+        _python_toolchain_build_file_content(
+            rctx.attr.toolchain_prefixes,
+            rctx.attr.toolchain_python_versions,
+            rctx.attr.toolchain_set_python_version_constraints,
+            rctx.attr.toolchain_user_repository_names,
+            rctx.attr._rules_python_workspace,
+        ),
+        executable = False,
+    )
+
+    (os, arch) = get_host_os_arch(rctx)
+    platform = get_host_platform(os, arch)
+    is_windows = (os == WINDOWS_NAME)
+    path = "python.exe" if is_windows else "bin/python3"
+
+    # Create a dict that is later used to create
+    # a symlink to a interpreter.
+    interpreter_labels = "".join([_line_for_hub_template.format(
+        name = name,
+        platform = platform,
+        path = path,
+    ) for name in rctx.attr.toolchain_user_repository_names])
+
+    rctx.file(
+        "interpreters.bzl",
+        _build_file_for_hub_template.format(
+            interpreter_labels = interpreter_labels,
+        ),
+        executable = False,
+    )
+
+hub_repo = repository_rule(
+    doc = """\
+This private rule create a repo with a BUILD file that contains a map of interpreter names
+and the labels to said interpreters. This map is used to by the interpreter hub extension.
+This rule also writes out the various toolchains for the different Python versions.
+""",
+    implementation = _hub_repo_impl,
+    attrs = {
+        "toolchain_prefixes": attr.string_list(
+            doc = "List prefixed for the toolchains",
+            mandatory = True,
+        ),
+        "toolchain_python_versions": attr.string_list(
+            doc = "List of Python versions for the toolchains",
+            mandatory = True,
+        ),
+        "toolchain_set_python_version_constraints": attr.string_list(
+            doc = "List of version contraints for the toolchains",
+            mandatory = True,
+        ),
+        "toolchain_user_repository_names": attr.string_list(
+            doc = "List of the user repo names for the toolchains",
+            mandatory = True,
+        ),
+        "_rules_python_workspace": attr.label(default = Label("//:does_not_matter_what_this_name_is")),
+    },
+)
diff --git a/python/extensions/python.bzl b/python/extensions/python.bzl
index cae1988..4732cfb 100644
--- a/python/extensions/python.bzl
+++ b/python/extensions/python.bzl
@@ -14,8 +14,28 @@
 
 "Python toolchain module extensions for use with bzlmod"
 
-load("@rules_python//python:repositories.bzl", "python_register_toolchains")
-load("@rules_python//python/extensions/private:interpreter_hub.bzl", "hub_repo")
+load("//python:repositories.bzl", "python_register_toolchains")
+load("//python/extensions/private:pythons_hub.bzl", "hub_repo")
+load("//python/private:toolchains_repo.bzl", "multi_toolchain_aliases")
+
+# This limit can be increased essentially arbitrarily, but doing so will cause a rebuild of all
+# targets using any of these toolchains due to the changed repository name.
+_MAX_NUM_TOOLCHAINS = 9999
+_TOOLCHAIN_INDEX_PAD_LENGTH = len(str(_MAX_NUM_TOOLCHAINS))
+
+def _toolchain_prefix(index, name):
+    """Prefixes the given name with the index, padded with zeros to ensure lexicographic sorting.
+
+    Examples:
+      _toolchain_prefix(   2, "foo") == "_0002_foo_"
+      _toolchain_prefix(2000, "foo") == "_2000_foo_"
+    """
+    return "_{}_{}_".format(_left_pad_zero(index, _TOOLCHAIN_INDEX_PAD_LENGTH), name)
+
+def _left_pad_zero(index, length):
+    if index < 0:
+        fail("index must be non-negative")
+    return ("0" * length + str(index))[-length:]
 
 # Printing a warning msg not debugging, so we have to disable
 # the buildifier check.
@@ -24,6 +44,8 @@ def _print_warn(msg):
     print("WARNING:", msg)
 
 def _python_register_toolchains(toolchain_attr, version_constraint):
+    """Calls python_register_toolchains and returns a struct used to collect the toolchains.
+    """
     python_register_toolchains(
         name = toolchain_attr.name,
         python_version = toolchain_attr.python_version,
@@ -31,24 +53,33 @@ def _python_register_toolchains(toolchain_attr, version_constraint):
         ignore_root_user_error = toolchain_attr.ignore_root_user_error,
         set_python_version_constraint = version_constraint,
     )
+    return struct(
+        python_version = toolchain_attr.python_version,
+        set_python_version_constraint = str(version_constraint),
+        name = toolchain_attr.name,
+    )
 
 def _python_impl(module_ctx):
-    # 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 = []
+    # Use to store all of the toolchains
+    toolchains = []
 
-    # Used to store the default toolchain name so we can pass it to the hub
-    default_toolchain_name = None
+    # Used to check if toolchains already exist
+    toolchain_names = []
 
     # Used to store toolchains that are in sub modules.
     sub_toolchains_map = {}
+    default_toolchain = None
+    python_versions = {}
 
     for mod in module_ctx.modules:
         for toolchain_attr in mod.tags.toolchain:
             # 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:
+                if toolchain_attr.name in toolchain_names:
+                    fail("""We found more than one toolchain that is named: {}.
+All toolchains must have an unique name.""".format(toolchain_attr.name))
+
                 toolchain_names.append(toolchain_attr.name)
 
                 # If we have the default version or we only have one toolchain
@@ -56,17 +87,27 @@ def _python_impl(module_ctx):
                 if toolchain_attr.is_default or len(mod.tags.toolchain) == 1:
                     # We have already found one default toolchain, and we can
                     # only have one.
-                    if default_toolchain_name != None:
+                    if default_toolchain != None:
                         fail("""We found more than one toolchain that is marked 
 as the default version.  Only set one toolchain with is_default set as 
 True. The toolchain is named: {}""".format(toolchain_attr.name))
 
-                    # We register the default toolchain.
-                    _python_register_toolchains(toolchain_attr, False)
-                    default_toolchain_name = toolchain_attr.name
-                else:
-                    #  Always register toolchains that are in the root module.
-                    _python_register_toolchains(toolchain_attr, version_constraint = True)
+                    # We store the default toolchain to have it
+                    # as the last toolchain added to toolchains
+                    default_toolchain = _python_register_toolchains(
+                        toolchain_attr,
+                        version_constraint = False,
+                    )
+                    python_versions[toolchain_attr.python_version] = toolchain_attr.name
+                    continue
+
+                toolchains.append(
+                    _python_register_toolchains(
+                        toolchain_attr,
+                        version_constraint = True,
+                    ),
+                )
+                python_versions[toolchain_attr.python_version] = toolchain_attr.name
             else:
                 # We add the toolchain to a map, and we later create the
                 # toolchain if the root module does not have a toolchain with
@@ -75,7 +116,7 @@ True. The toolchain is named: {}""".format(toolchain_attr.name))
                 sub_toolchains_map[toolchain_attr.name] = toolchain_attr
 
     # We did not find a default toolchain so we fail.
-    if default_toolchain_name == None:
+    if default_toolchain == None:
         fail("""Unable to find a default toolchain in the root module.  
 Please define a toolchain that has is_version set to True.""")
 
@@ -88,14 +129,39 @@ Please define a toolchain that has is_version set to True.""")
  module has a toolchain of the same name.""".format(toolchain_attr.name))
             continue
         toolchain_names.append(name)
-        _python_register_toolchains(toolchain_attr, True)
-
-    # Create the hub for the interpreters and the
-    # the default toolchain.
+        toolchains.append(
+            _python_register_toolchains(
+                toolchain_attr,
+                version_constraint = True,
+            ),
+        )
+        python_versions[toolchain_attr.python_version] = toolchain_attr.name
+
+    # The last toolchain in the BUILD file is set as the default
+    # toolchain. We need the default last.
+    toolchains.append(default_toolchain)
+
+    if len(toolchains) > _MAX_NUM_TOOLCHAINS:
+        fail("more than {} python versions are not supported".format(_MAX_NUM_TOOLCHAINS))
+
+    # Create the pythons_hub repo for the interpreter meta data and the
+    # the various toolchains.
     hub_repo(
         name = "pythons_hub",
-        toolchains = toolchain_names,
-        default_toolchain = default_toolchain_name,
+        toolchain_prefixes = [
+            _toolchain_prefix(index, toolchain.name)
+            for index, toolchain in enumerate(toolchains)
+        ],
+        toolchain_python_versions = [t.python_version for t in toolchains],
+        toolchain_set_python_version_constraints = [t.set_python_version_constraint for t in toolchains],
+        toolchain_user_repository_names = [t.name for t in toolchains],
+    )
+
+    # This is require in order to support multiple version py_test
+    # and py_binary
+    multi_toolchain_aliases(
+        name = "python_aliases",
+        python_versions = python_versions,
     )
 
 python = module_extension(
@@ -142,8 +208,14 @@ is set as the default toolchain.
                     mandatory = False,
                     doc = "Whether the toolchain is the default version",
                 ),
-                "name": attr.string(mandatory = True),
-                "python_version": attr.string(mandatory = True),
+                "name": attr.string(
+                    mandatory = True,
+                    doc = "Name of the toolchain",
+                ),
+                "python_version": attr.string(
+                    mandatory = True,
+                    doc = "The Python version that we are creating the toolchain for.",
+                ),
             },
         ),
     },
diff --git a/python/private/toolchains_repo.bzl b/python/private/toolchains_repo.bzl
index 9bed73e..b5ac81a 100644
--- a/python/private/toolchains_repo.bzl
+++ b/python/private/toolchains_repo.bzl
@@ -35,28 +35,37 @@ def get_repository_name(repository_workspace):
     dummy_label = "//:_"
     return str(repository_workspace.relative(dummy_label))[:-len(dummy_label)] or "@"
 
-def _toolchains_repo_impl(rctx):
-    python_version_constraint = "{rules_python}//python/config_settings:is_python_{python_version}".format(
-        rules_python = get_repository_name(rctx.attr._rules_python_workspace),
-        python_version = rctx.attr.python_version,
-    )
+def python_toolchain_build_file_content(
+        prefix,
+        python_version,
+        set_python_version_constraint,
+        user_repository_name,
+        rules_python):
+    """Creates the content for toolchain definitions for a build file.
 
-    build_content = """\
-# Generated by python/private/toolchains_repo.bzl
-#
-# These can be registered in the workspace file or passed to --extra_toolchains
-# flag. By default all these toolchains are registered by the
-# python_register_toolchains macro so you don't normally need to interact with
-# these targets.
+    Args:
+        prefix: Python toolchain name prefixes
+        python_version: Python versions for the toolchains
+        set_python_version_constraint: string "True" or "False"
+        user_repository_name: names for the user repos
+        rules_python: rules_python label
 
-"""
+    Returns:
+        build_content: Text containing toolchain definitions
+    """
+
+    python_version_constraint = "{rules_python}//python/config_settings:is_python_{python_version}".format(
+        rules_python = rules_python,
+        python_version = python_version,
+    )
 
-    for [platform, meta] in PLATFORMS.items():
-        build_content += """\
-# Bazel selects this toolchain to get a Python interpreter
-# for executing build actions.
+    # We create a list of toolchain content from iterating over
+    # the enumeration of PLATFORMS.  We enumerate PLATFORMS in
+    # order to get us an index to increment the increment.
+    return "".join([
+        """
 toolchain(
-    name = "{platform}_toolchain",
+    name = "{prefix}{platform}_toolchain",
     target_compatible_with = {compatible_with},
     target_settings = ["{python_version_constraint}"] if {set_python_version_constraint} else [],
     toolchain = "@{user_repository_name}_{platform}//:python_runtimes",
@@ -64,14 +73,42 @@ toolchain(
 )
 """.format(
             compatible_with = meta.compatible_with,
-            name = rctx.attr.name,
             platform = platform,
             python_version_constraint = python_version_constraint,
-            set_python_version_constraint = rctx.attr.set_python_version_constraint,
-            user_repository_name = rctx.attr.user_repository_name,
+            # We have to use a String value here because bzlmod is passing in a
+            # string as we cannot have list of bools in build rule attribues.
+            # This if statement does not appear to work unless it is in the
+            # toolchain file.
+            set_python_version_constraint = True if set_python_version_constraint == "True" else False,
+            user_repository_name = user_repository_name,
+            prefix = prefix,
         )
+        for platform, meta in PLATFORMS.items()
+    ])
+
+def _toolchains_repo_impl(rctx):
+    build_content = """\
+# Generated by python/private/toolchains_repo.bzl
+#
+# These can be registered in the workspace file or passed to --extra_toolchains
+# flag. By default all these toolchains are registered by the
+# python_register_toolchains macro so you don't normally need to interact with
+# these targets.
+
+"""
+
+    # Get the repository name
+    rules_python = get_repository_name(rctx.attr._rules_python_workspace)
+
+    toolchains = python_toolchain_build_file_content(
+        prefix = "",
+        python_version = rctx.attr.python_version,
+        set_python_version_constraint = str(rctx.attr.set_python_version_constraint),
+        user_repository_name = rctx.attr.user_repository_name,
+        rules_python = rules_python,
+    )
 
-    rctx.file("BUILD.bazel", build_content)
+    rctx.file("BUILD.bazel", build_content + toolchains)
 
 toolchains_repo = repository_rule(
     _toolchains_repo_impl,
diff --git a/python/repositories.bzl b/python/repositories.bzl
index 4f36b12..e841e28 100644
--- a/python/repositories.bzl
+++ b/python/repositories.bzl
@@ -564,16 +564,20 @@ def python_register_toolchains(
                 platform = platform,
             ))
 
-    toolchains_repo(
-        name = toolchain_repo_name,
+    toolchain_aliases(
+        name = name,
         python_version = python_version,
-        set_python_version_constraint = set_python_version_constraint,
         user_repository_name = name,
     )
 
-    toolchain_aliases(
-        name = name,
+    # in bzlmod we write out our own toolchain repos
+    if bzlmod:
+        return
+
+    toolchains_repo(
+        name = toolchain_repo_name,
         python_version = python_version,
+        set_python_version_constraint = set_python_version_constraint,
         user_repository_name = name,
     )