feat: allow specifying arbitrary constraints for local toolchains (#2829)

This adds the ability for local toolchains to have arbitrary constraints
set on them.

This allows accomplishing two goals:
1. Makes it easier to enable/disable them on the command line, instead
of having them
entirely override an existing config and having to comment/uncomment the
   MODULE.bazel file sections.
2. Allows configuring them so that the repository is never initialized,
which avoids
the repository from being initialized during toolchain resolution, even
if it will
   never match because of (1).

Author: rickeylev

.bazelci/presubmit.yml

@@ -51,9 +51,11 @@ buildifier:
   test_flags:
     - "--noenable_bzlmod"
     - "--enable_workspace"
+    - "--test_tag_filters=-integration-test"
   build_flags:
     - "--noenable_bzlmod"
     - "--enable_workspace"
+    - "--build_tag_filters=-integration-test"
   bazel: 7.x
 .common_bazelinbazel_config: &common_bazelinbazel_config
     build_flags:

CHANGELOG.md

@@ -84,6 +84,8 @@ END_UNRELEASED_TEMPLATE
 * Repo utilities `execute_unchecked`, `execute_checked`, and `execute_checked_stdout` now
   support `log_stdout` and `log_stderr` keyword arg booleans. When these are `True`
   (the default), the subprocess's stdout/stderr will be logged.
+* (toolchains) Local toolchains can be activated with custom flags. See
+  [Conditionally using local toolchains] docs for how to configure.
 
 {#v0-0-0-removed}
 ### Removed

docs/toolchains.md

@@ -377,15 +377,14 @@ local_runtime_repo(
 local_runtime_toolchains_repo(
     name = "local_toolchains",
     runtimes = ["local_python3"],
+    # TIP: The `target_settings` arg can be used to activate them based on
+    # command line flags; see docs below.
 )
 
 # Step 3: Register the toolchains
 register_toolchains("@local_toolchains//:all", dev_dependency = True)
 ```
 
-Note that `register_toolchains` will insert the local toolchain earlier in the
-toolchain ordering, so it will take precedence over other registered toolchains.
-
 :::{important}
 Be sure to set `dev_dependency = True`. Using a local toolchain only makes sense
 for the root module.
@@ -397,6 +396,72 @@ downstream modules.
 
 Multiple runtimes and/or toolchains can be defined, which allows for multiple
 Python versions and/or platforms to be configured in a single `MODULE.bazel`.
+Note that `register_toolchains` will insert the local toolchain earlier in the
+toolchain ordering, so it will take precedence over other registered toolchains.
+To better control when the toolchain is used, see [Conditionally using local
+toolchains]
+
+### Conditionally using local toolchains
+
+By default, a local toolchain has few constraints and is early in the toolchain
+ordering, which means it will usually be used no matter what. This can be
+problematic for CI (where it shouldn't be used), expensive for CI (CI must
+initialize/download the repository to determine its Python version), and
+annoying for iterative development (enabling/disabling it requires modifying
+MODULE.bazel).
+
+These behaviors can be mitigated, but it requires additional configuration
+to avoid triggering the local toolchain repository to initialize (i.e. run
+local commands and perform downloads).
+
+The two settings to change are
+{obj}`local_runtime_toolchains_repo.target_compatible_with` and
+{obj}`local_runtime_toolchains_repo.target_settings`, which control how Bazel
+decides if a toolchain should match. By default, they point to targets *within*
+the local runtime repository (trigger repo initialization). We have to override
+them to *not* reference the local runtime repository at all.
+
+In the example below, we reconfigure the local toolchains so they are only
+activated if the custom flag `--//:py=local` is set and the target platform
+matches the Bazel host platform. The net effect is CI won't use the local
+toolchain (nor initialize its repository), and developers can easily
+enable/disable the local toolchain with a command line flag.
+
+```
+# File: MODULE.bazel
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
+
+local_runtime_toolchains_repo(
+    name = "local_toolchains",
+    runtimes = ["local_python3"],
+    target_compatible_with = {
+        "local_python3": ["HOST_CONSTRAINTS"],
+    },
+    target_settings = {
+        "local_python3": ["@//:is_py_local"]
+    }
+)
+
+# File: BUILD.bazel
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+
+config_setting(
+    name = "is_py_local",
+    flag_values = {":py": "local"},
+)
+
+string_flag(
+    name = "py",
+    build_setting_default = "",
+)
+```
+
+:::{tip}
+Easily switching between *multiple* local toolchains can be accomplished by
+adding additional `:is_py_X` targets and setting `--//:py` to match.
+to easily switch between different local toolchains.
+:::
+
 
 ## Runtime environment toolchain
 
@@ -425,7 +490,7 @@ locally installed Python.
 ### Autodetecting toolchain
 
 The autodetecting toolchain is a deprecated toolchain that is built into Bazel.
-It's name is a bit misleading: it doesn't autodetect anything. All it does is
+**It's name is a bit misleading: it doesn't autodetect anything**. All it does is
 use `python3` from the environment a binary runs within. This provides extremely
 limited functionality to the rules (at build time, nothing is knowable about
 the Python runtime).

python/private/local_runtime_toolchains_repo.bzl

@@ -26,6 +26,9 @@ define_local_toolchain_suites(
     name = "toolchains",
     version_aware_repo_names = {version_aware_names},
     version_unaware_repo_names = {version_unaware_names},
+    repo_exec_compatible_with = {repo_exec_compatible_with},
+    repo_target_compatible_with = {repo_target_compatible_with},
+    repo_target_settings = {repo_target_settings},
 )
 """
 
@@ -39,6 +42,9 @@ def _local_runtime_toolchains_repo(rctx):
 
     rctx.file("BUILD.bazel", _TOOLCHAIN_TEMPLATE.format(
         version_aware_names = render.list(rctx.attr.runtimes),
+        repo_target_settings = render.string_list_dict(rctx.attr.target_settings),
+        repo_target_compatible_with = render.string_list_dict(rctx.attr.target_compatible_with),
+        repo_exec_compatible_with = render.string_list_dict(rctx.attr.exec_compatible_with),
         version_unaware_names = render.list(rctx.attr.default_runtimes or rctx.attr.runtimes),
     ))
 
@@ -62,8 +68,36 @@ These will be defined as *version-unaware* toolchains. This means they will
 match any Python version. As such, they are registered after the version-aware
 toolchains defined by the `runtimes` attribute.
 
+If not set, then the `runtimes` values will be used.
+
 Note that order matters: it determines the toolchain priority within the
 package.
+""",
+        ),
+        "exec_compatible_with": attr.string_list_dict(
+            doc = """
+Constraints that must be satisfied by an exec platform for a toolchain to be used.
+
+This is a `dict[str, list[str]]`, where the keys are repo names from the
+`runtimes` or `default_runtimes` args, and the values are constraint
+target labels (e.g. OS, CPU, etc).
+
+:::{note}
+Specify `@//foo:bar`, not simply `//foo:bar` or `:bar`. The additional `@` is
+needed because the strings are evaluated in a different context than where
+they originate.
+:::
+
+The list of settings become the {obj}`toolchain.exec_compatible_with` value for
+each respective repo.
+
+This allows a local toolchain to only be used if certain exec platform
+conditions are met, typically values from `@platforms`.
+
+See the [Local toolchains] docs for examples and further information.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
 """,
         ),
         "runtimes": attr.string_list(
@@ -76,6 +110,81 @@ are registered before `default_runtimes`.
 
 Note that order matters: it determines the toolchain priority within the
 package.
+""",
+        ),
+        "target_compatible_with": attr.string_list_dict(
+            doc = """
+Constraints that must be satisfied for a toolchain to be used.
+
+
+This is a `dict[str, list[str]]`, where the keys are repo names from the
+`runtimes` or `default_runtimes` args, and the values are constraint
+target labels (e.g. OS, CPU, etc), or the special string `"HOST_CONSTRAINTS"`
+(which will be replaced with the current Bazel hosts's constraints).
+
+If a repo's entry is missing or empty, it defaults to the supported OS the
+underlying runtime repository detects as compatible.
+
+:::{note}
+Specify `@//foo:bar`, not simply `//foo:bar` or `:bar`. The additional `@` is
+needed because the strings are evaluated in a different context than where
+they originate.
+:::
+
+The list of settings **becomes the** the {obj}`toolchain.target_compatible_with`
+value for each respective repo; i.e. they _replace_ the auto-detected values
+the local runtime itself computes.
+
+This allows a local toolchain to only be used if certain target platform
+conditions are met, typically values from `@platforms`.
+
+See the [Local toolchains] docs for examples and further information.
+
+:::{seealso}
+The `target_settings` attribute, which handles `config_setting` values,
+instead of constraints.
+:::
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+        ),
+        "target_settings": attr.string_list_dict(
+            doc = """
+Config settings that must be satisfied for a toolchain to be used.
+
+This is a `dict[str, list[str]]`, where the keys are repo names from the
+`runtimes` or `default_runtimes` args, and the values are {obj}`config_setting()`
+target labels.
+
+If a repo's entry is missing or empty, it will default to
+`@<repo>//:is_match_python_version` (for repos in `runtimes`) or an empty list
+(for repos in `default_runtimes`).
+
+:::{note}
+Specify `@//foo:bar`, not simply `//foo:bar` or `:bar`. The additional `@` is
+needed because the strings are evaluated in a different context than where
+they originate.
+:::
+
+The list of settings will be applied atop of any of the local runtime's
+settings that are used for {obj}`toolchain.target_settings`. i.e. they are
+evaluated first and guard the checking of the local runtime's auto-detected
+conditions.
+
+This allows a local toolchain to only be used if certain flags or
+config setting conditions are met. Such conditions can include user-defined
+flags, platform constraints, etc.
+
+See the [Local toolchains] docs for examples and further information.
+
+:::{seealso}
+The `target_compatible_with` attribute, which handles *constraint* values,
+instead of `config_settings`.
+:::
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
 """,
         ),
         "_rule_name": attr.string(default = "local_toolchains_repo"),

python/private/py_toolchain_suite.bzl

@@ -15,6 +15,7 @@
 """Create the toolchain defs in a BUILD.bazel file."""
 
 load("@bazel_skylib//lib:selects.bzl", "selects")
+load("@platforms//host:constraints.bzl", "HOST_CONSTRAINTS")
 load(":text_util.bzl", "render")
 load(
     ":toolchain_types.bzl",
@@ -95,9 +96,15 @@ def py_toolchain_suite(
         runtime_repo_name = user_repository_name,
         target_settings = target_settings,
         target_compatible_with = target_compatible_with,
+        exec_compatible_with = [],
     )
 
-def _internal_toolchain_suite(prefix, runtime_repo_name, target_compatible_with, target_settings):
+def _internal_toolchain_suite(
+        prefix,
+        runtime_repo_name,
+        target_compatible_with,
+        target_settings,
+        exec_compatible_with):
     native.toolchain(
         name = "{prefix}_toolchain".format(prefix = prefix),
         toolchain = "@{runtime_repo_name}//:python_runtimes".format(
@@ -106,6 +113,7 @@ def _internal_toolchain_suite(prefix, runtime_repo_name, target_compatible_with,
         toolchain_type = TARGET_TOOLCHAIN_TYPE,
         target_settings = target_settings,
         target_compatible_with = target_compatible_with,
+        exec_compatible_with = exec_compatible_with,
     )
 
     native.toolchain(
@@ -116,6 +124,7 @@ def _internal_toolchain_suite(prefix, runtime_repo_name, target_compatible_with,
         toolchain_type = PY_CC_TOOLCHAIN_TYPE,
         target_settings = target_settings,
         target_compatible_with = target_compatible_with,
+        exec_compatible_with = exec_compatible_with,
     )
 
     native.toolchain(
@@ -142,7 +151,13 @@ def _internal_toolchain_suite(prefix, runtime_repo_name, target_compatible_with,
     # call in python/repositories.bzl. Bzlmod doesn't need anything; it will
     # register `:all`.
 
-def define_local_toolchain_suites(name, version_aware_repo_names, version_unaware_repo_names):
+def define_local_toolchain_suites(
+        name,
+        version_aware_repo_names,
+        version_unaware_repo_names,
+        repo_exec_compatible_with,
+        repo_target_compatible_with,
+        repo_target_settings):
     """Define toolchains for `local_runtime_repo` backed toolchains.
 
     This generates `toolchain` targets that can be registered using `:all`. The
@@ -156,24 +171,60 @@ def define_local_toolchain_suites(name, version_aware_repo_names, version_unawar
             version-aware toolchains defined.
         version_unaware_repo_names: `list[str]` of the repo names that will have
             version-unaware toolchains defined.
+        repo_target_settings: {type}`dict[str, list[str]]` mapping of repo names
+            to string labels that are added to the `target_settings` for the
+            respective repo's toolchain.
+        repo_target_compatible_with: {type}`dict[str, list[str]]` mapping of repo names
+            to string labels that are added to the `target_compatible_with` for
+            the respective repo's toolchain.
+        repo_exec_compatible_with: {type}`dict[str, list[str]]` mapping of repo names
+            to string labels that are added to the `exec_compatible_with` for
+            the respective repo's toolchain.
     """
+
     i = 0
     for i, repo in enumerate(version_aware_repo_names, start = i):
-        prefix = render.left_pad_zero(i, 4)
+        target_settings = ["@{}//:is_matching_python_version".format(repo)]
+
+        if repo_target_settings.get(repo):
+            selects.config_setting_group(
+                name = "_{}_user_guard".format(repo),
+                match_all = repo_target_settings.get(repo, []) + target_settings,
+            )
+            target_settings = ["_{}_user_guard".format(repo)]
         _internal_toolchain_suite(
-            prefix = prefix,
+            prefix = render.left_pad_zero(i, 4),
             runtime_repo_name = repo,
-            target_compatible_with = ["@{}//:os".format(repo)],
-            target_settings = ["@{}//:is_matching_python_version".format(repo)],
+            target_compatible_with = _get_local_toolchain_target_compatible_with(
+                repo,
+                repo_target_compatible_with,
+            ),
+            target_settings = target_settings,
+            exec_compatible_with = repo_exec_compatible_with.get(repo, []),
         )
 
     # The version unaware entries must go last because they will match any Python
     # version.
     for i, repo in enumerate(version_unaware_repo_names, start = i + 1):
-        prefix = render.left_pad_zero(i, 4)
         _internal_toolchain_suite(
-            prefix = prefix,
+            prefix = render.left_pad_zero(i, 4) + "_default",
             runtime_repo_name = repo,
-            target_settings = [],
-            target_compatible_with = ["@{}//:os".format(repo)],
+            target_compatible_with = _get_local_toolchain_target_compatible_with(
+                repo,
+                repo_target_compatible_with,
+            ),
+            # We don't call _get_local_toolchain_target_settings because that
+            # will add the version matching condition by default.
+            target_settings = repo_target_settings.get(repo, []),
+            exec_compatible_with = repo_exec_compatible_with.get(repo, []),
         )
+
+def _get_local_toolchain_target_compatible_with(repo, repo_target_compatible_with):
+    if repo in repo_target_compatible_with:
+        target_compatible_with = repo_target_compatible_with[repo]
+        if "HOST_CONSTRAINTS" in target_compatible_with:
+            target_compatible_with.remove("HOST_CONSTRAINTS")
+            target_compatible_with.extend(HOST_CONSTRAINTS)
+    else:
+        target_compatible_with = ["@{}//:os".format(repo)]
+    return target_compatible_with