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

We do this work for two reasons.

First, we must support Module dependencies and sub-modules using
`python.toolchain`. There are already two known instances of sub-modules
setting up a Python toolchain and colliding with another module
(nanobind and rules_testing both run into this).

Second, the upcoming multi-version support is going to work by having
each `python.toolchain()` call register its particular version with the
extra toolchain constraint. This also helps unify the version-aware and
non-version-aware code paths (the non-version aware paths are just
version-aware with a single version registered as the default)

This commit implements various business logic in the toolchain class.

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.

See #1229 for more information

Author: chrislovecnm

examples/bzlmod/.bazelrc

@@ -1,3 +1,10 @@
 common --experimental_enable_bzlmod
 
 coverage --java_runtime_version=remotejdk_11
+
+test --test_output=errors --enable_runfiles
+
+# Windows requires these for multi-python support:
+build --enable_runfiles
+
+startup --windows_enable_symlinks

examples/bzlmod/BUILD.bazel

@@ -1,9 +1,21 @@
+# Load various rules so that we can have bazel download
+# various rulesets and dependencies.
+# The `load` statement imports the symbol for the rule, in the defined
+# ruleset. When the symbol is loaded you can use the rule.
+
+# The names @pip and @python_39 are values that are repository
+# names. Those names are defined in the MODULES.bazel file.
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load("@pip//:requirements.bzl", "all_requirements", "all_whl_requirements", "requirement")
-load("@python3_9//:defs.bzl", py_test_with_transition = "py_test")
+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")
 
+# This stanza calls a rule that generates targets for managing pip dependencies
+# with pip-compile.
 compile_pip_requirements(
     name = "requirements",
     extra_args = ["--allow-unsafe"],
@@ -12,6 +24,11 @@ compile_pip_requirements(
     requirements_windows = "requirements_windows.txt",
 )
 
+# The rules below are language specific rules defined in
+# rules_python. See
+# https://bazel.build/reference/be/python
+
+# see https://bazel.build/reference/be/python#py_library
 py_library(
     name = "lib",
     srcs = ["lib.py"],
@@ -22,6 +39,7 @@ py_library(
     ],
 )
 
+# see https://bazel.build/reference/be/python#py_binary
 py_binary(
     name = "bzlmod",
     srcs = ["__main__.py"],
@@ -32,6 +50,23 @@ 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",
     srcs = ["test.py"],
@@ -46,6 +81,11 @@ py_test_with_transition(
     deps = [":lib"],
 )
 
+# This example is also used for integration tests within
+# rules_python.  We are using
+# https://github.com/bazelbuild/bazel-skylib
+# to run some of the tests.
+# See: https://github.com/bazelbuild/bazel-skylib/blob/main/docs/build_test_doc.md
 build_test(
     name = "all_wheels",
     targets = all_whl_requirements,

examples/bzlmod/MODULE.bazel

@@ -11,32 +11,89 @@ local_path_override(
     path = "../..",
 )
 
+# This name is passed into python.toolchain and it's use_repo statement.
+# 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.
+# You can set different Python versions in this block.
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
-    name = "python3_9",
+    # This name is used in the various use_repo statements
+    # below, and in the local extension that is in this
+    # example.
+    name = PYTHON_NAME_39,
     configure_coverage_tool = True,
+    # Only set when you have mulitple toolchain versions.
+    is_default = True,
     python_version = "3.9",
 )
-use_repo(python, "python3_9")
-use_repo(python, "python3_9_toolchains")
 
+# We are also using a second version of Python in this project.
+# Typically you will only need a single version of Python, but
+# If you need a different vesion we support more than one.
+# Note: we do not supporting using multiple pip extensions, this is
+# work in progress.
+python.toolchain(
+    name = PYTHON_NAME_310,
+    configure_coverage_tool = True,
+    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(
-    "@python3_9_toolchains//:all",
+    "@{}//:all".format(PYTHON_39_TOOLCHAINS),
+    "@{}//:all".format(PYTHON_310_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.
 interpreter = use_extension("@rules_python//python/extensions:interpreter.bzl", "interpreter")
 interpreter.install(
-    name = "interpreter_python3_9",
-    python_name = "python3_9",
+    name = INTERPRETER_NAME_39,
+    python_name = PYTHON_NAME_39,
+)
+
+# This second call is only needed if you are using mulitple different
+# Python versions/interpreters.
+interpreter.install(
+    name = INTERPRETER_NAME_310,
+    python_name = PYTHON_NAME_310,
 )
-use_repo(interpreter, "interpreter_python3_9")
+use_repo(interpreter, INTERPRETER_NAME_39, INTERPRETER_NAME_310)
 
 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",
+    python_interpreter_target = "@{}//:python".format(INTERPRETER_NAME_39),
     requirements_lock = "//:requirements_lock.txt",
     requirements_windows = "//:requirements_windows.txt",
 )

examples/bzlmod/__main__.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from lib import main
+import sys
 
 if __name__ == "__main__":
     print(main([["A", 1], ["B", 2]]))
+    print(sys.version)

examples/bzlmod/other_module/MODULE.bazel

@@ -2,4 +2,56 @@ module(
     name = "other_module",
 )
 
+# This module is using the same version of rules_python
+# that the parent module uses.
 bazel_dep(name = "rules_python", version = "")
+
+# It is not best practice to use a python.toolchian in
+# a submodule.  This code only exists to test that
+# we support doing this. This code is only for rules_python
+# 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
+    # below, and in the local extension that is in this
+    # example.
+    name = PYTHON_NAME_39,
+    configure_coverage_tool = True,
+    python_version = "3.9",
+)
+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_311,
+    configure_coverage_tool = True,
+    # In a submodule this is ignored
+    is_default = True,
+    python_version = "3.11",
+)
+
+# created by the above python.toolchain calls.
+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),
+)

examples/bzlmod/other_module/other_module/pkg/BUILD.bazel

@@ -1,3 +1,4 @@
+load("@python_311//:defs.bzl", py_binary_311 = "py_binary")
 load("@rules_python//python:defs.bzl", "py_library")
 
 py_library(
@@ -8,4 +9,15 @@ py_library(
     deps = ["@rules_python//python/runfiles"],
 )
 
+# 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_311(
+    name = "lib_311",
+    srcs = ["lib.py"],
+    data = ["data/data.txt"],
+    visibility = ["//visibility:public"],
+    deps = ["@rules_python//python/runfiles"],
+)
+
 exports_files(["data/data.txt"])

examples/bzlmod/runfiles/BUILD.bazel

@@ -1,6 +1,5 @@
 load("@rules_python//python:defs.bzl", "py_test")
 
-# gazelle:ignore
 py_test(
     name = "runfiles_test",
     srcs = ["runfiles_test.py"],

examples/bzlmod_build_file_generation/BUILD.bazel

@@ -7,7 +7,7 @@
 # requirements.
 load("@bazel_gazelle//:def.bzl", "gazelle")
 load("@pip//:requirements.bzl", "all_whl_requirements")
-load("@python3//:defs.bzl", py_test_with_transition = "py_test")
+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")

examples/bzlmod_build_file_generation/MODULE.bazel

@@ -46,7 +46,11 @@ python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 
 # This name is passed into python.toolchain and it's use_repo statement.
 # We also use the same name for python.host_python_interpreter.
-PYTHON_NAME = "python3"
+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.
@@ -56,30 +60,35 @@ python.toolchain(
     # 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, "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.
 # It creates a symlink to the binary, and we pass the label to the following
 # pip.parse call.
 interpreter = use_extension("@rules_python//python/extensions:interpreter.bzl", "interpreter")
 interpreter.install(
-    name = "interpreter_python3",
+    name = INTERPRETER_NAME,
     python_name = PYTHON_NAME,
 )
-use_repo(interpreter, "interpreter_python3")
+use_repo(interpreter, INTERPRETER_NAME)
 
 # Use the extension, pip.parse, to call the `pip_repository` rule that invokes
 # `pip`, with `incremental` set. The pip call accepts a locked/compiled
@@ -102,7 +111,7 @@ pip.parse(
     # is used for both resolving dependencies and running tests/binaries.
     # If this isn't specified, then you'll get whatever is locally installed
     # on your system.
-    python_interpreter_target = "@interpreter_python3//:python",
+    python_interpreter_target = "@{}//:python".format(INTERPRETER_NAME),
     requirements_lock = "//:requirements_lock.txt",
     requirements_windows = "//:requirements_windows.txt",
 )

python/extensions/private/interpreter_hub.bzl

@@ -19,8 +19,9 @@ load("//python/private:toolchains_repo.bzl", "get_host_os_arch", "get_host_platf
 
 _build_file_for_hub_template = """
 INTERPRETER_LABELS = {{
-{lines}
+{interpreter_labels}
 }}
+DEFAULT_TOOLCHAIN_NAME = "{default}"
 """
 
 _line_for_hub_template = """\
@@ -35,13 +36,19 @@ def _hub_repo_impl(rctx):
     is_windows = (os == WINDOWS_NAME)
     path = "python.exe" if is_windows else "bin/python3"
 
-    lines = "\n".join([_line_for_hub_template.format(
+    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(lines = lines))
+    rctx.file(
+        "interpreters.bzl",
+        _build_file_for_hub_template.format(
+            interpreter_labels = interpreter_labels,
+            default = rctx.attr.default_toolchain,
+        ),
+    )
 
 hub_repo = repository_rule(
     doc = """\
@@ -50,6 +57,10 @@ and the labels to said interpreters. This map is used to by the interpreter hub
 """,
     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,