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

chrislovecnm · chrislovecnm · commit fa0dc381373f · 2023-05-16T18:40:20.000-06:00
We do this work for two reasons. First, we must support Module dependencies and sub-modules using` python.toolchain`. Second, we needed this commit in order to support 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 is the same naming convention that python_register_multi_toolchains uses. See: https://github.com/bazelbuild/rules_python/blob/16126d0ebfee074a3a8fe216b20cc19e1b3603c1/python/repositories.bzl#L569 This commit also updates the two bzlmod examples.
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
@@ -10,24 +10,35 @@ 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")
+interpreter.install(
+    name = "interpreter_python_3_9",
+    python_name = PYTHON_NAME,
+)
+use_repo(interpreter, "interpreter_python_3_9")
+
+
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 pip.parse(
     name = "pip",
     requirements_lock = "//:requirements_lock.txt",
     requirements_windows = "//:requirements_windows.txt",
+    python_interpreter_target = "@interpreter_python_3_9//:python",
 )
 use_repo(pip, "pip")
 
diff --git a/examples/bzlmod_build_file_generation/MODULE.bazel b/examples/bzlmod_build_file_generation/MODULE.bazel
@@ -47,6 +47,7 @@ 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_TOOLCHAINS = PYTHON_NAME + "_toolchains"
 
 # We next initialize the python toolchain using the extension.
 # You can set different Python versions in this block.
@@ -63,12 +64,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.
diff --git a/python/extensions/interpreter.bzl b/python/extensions/interpreter.bzl
@@ -51,9 +51,11 @@ in the MODULES.bazel file.
 def _interpreter_repo_impl(rctx):
     rctx.file("BUILD.bazel", "")
 
-    actual_interpreter_label = INTERPRETER_LABELS.get(rctx.attr.python_name)
+    name = rctx.attr.python_name
+    actual_interpreter_label = INTERPRETER_LABELS.get(name)
+
     if actual_interpreter_label == None:
-        fail("Unable to find interpreter with name {}".format(rctx.attr.python_name))
+        fail("Unable to find interpreter with name {}".format(name))
 
     rctx.symlink(actual_interpreter_label, "python")
 
diff --git a/python/extensions/python.bzl b/python/extensions/python.bzl
@@ -17,25 +17,40 @@
 load("@rules_python//python:repositories.bzl", "python_register_toolchains")
 load("@rules_python//python/extensions/private:interpreter_hub.bzl", "hub_repo")
 
+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):
+    # We collect all of the toolchain names to create
+    # the INTERPRETER_LABELS map.  This is used
+    # by interpreter_extensions.bzl
     toolchains = []
+    default_toolchain = None
     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)
 
+            # We register the default toolchain last
+            # And if we only have one toolchain we register it as the default
+            # TODO see if we are in the root module
+            if toolchain_attr.default_version or len(mod.tags.toolchain) == 1:
+                default_toolchain = toolchain_attr
+                continue
+
+            # We are suffixing with the version in order to handle multiple
+            # toolchains.
+            _python_register_toolchains(toolchain_attr, True)
+
+    if default_toolchain:
+        _python_register_toolchains(default_toolchain, False)
+
+    # create the hub for the different interpreter versions
     hub_repo(
         name = "pythons_hub",
         toolchains = toolchains,
@@ -51,6 +66,10 @@ python = module_extension(
                     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.",
diff --git a/python/repositories.bzl b/python/repositories.bzl
@@ -457,7 +457,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.
 
@@ -480,9 +479,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:
