fix(toolchains): correctly order the toolchains (bazel-contrib#2735)

aignas · web-flow · commit 3d98aeea9c70 · 2025-04-05T13:17:28.000Z
Since toolchain matching is done by matching the first target that matches target settings, the `minor_mapping` config setting is special, because e.g. all `3.11.X` toolchains match the `python_version = "3.11"` setting. This just reshuffles the list so that we have toolchains that are in the `minor_mapping` before the rest. At the same time remove the workaround from the `lock.bzl` where the bug was initially discovered. Fixes bazel-contrib#2685
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -78,6 +78,9 @@ Unreleased changes template.
 * (toolchains) Do not try to run `chmod` when downloading non-windows hermetic toolchain
   repositories on Windows. Fixes
   [#2660](https://github.com/bazel-contrib/rules_python/issues/2660).
+* (toolchains) The toolchain matching is has been fixed when writing
+  transitions transitioning on the `python_version` flag.
+  Fixes [#2685](https://github.com/bazel-contrib/rules_python/issues/2685).
 
 {#v0-0-0-added}
 ### Added
diff --git a/python/private/python.bzl b/python/private/python.bzl
@@ -243,10 +243,25 @@ def parse_modules(*, module_ctx, _fail = fail):
     if len(toolchains) > _MAX_NUM_TOOLCHAINS:
         fail("more than {} python versions are not supported".format(_MAX_NUM_TOOLCHAINS))
 
+    # sort the toolchains so that the toolchain versions that are in the
+    # `minor_mapping` are coming first. This ensures that `python_version =
+    # "3.X"` transitions work as expected.
+    minor_version_toolchains = []
+    other_toolchains = []
+    minor_mapping = list(config.minor_mapping.values())
+    for t in toolchains:
+        # FIXME @aignas 2025-04-04: How can we unit test that this ordering is
+        # consistent with what would actually work?
+        if config.minor_mapping.get(t.python_version, t.python_version) in minor_mapping:
+            minor_version_toolchains.append(t)
+        else:
+            other_toolchains.append(t)
+    toolchains = minor_version_toolchains + other_toolchains
+
     return struct(
         config = config,
         debug_info = debug_info,
-        default_python_version = toolchains[-1].python_version,
+        default_python_version = default_toolchain.python_version,
         toolchains = [
             struct(
                 python_version = t.python_version,
diff --git a/python/uv/private/BUILD.bazel b/python/uv/private/BUILD.bazel
@@ -43,10 +43,8 @@ bzl_library(
         ":toolchain_types_bzl",
         "//python:py_binary_bzl",
         "//python/private:bzlmod_enabled_bzl",
-        "//python/private:full_version_bzl",
         "//python/private:toolchain_types_bzl",
         "@bazel_skylib//lib:shell",
-        "@pythons_hub//:versions_bzl",
     ],
 )
 
diff --git a/python/uv/private/lock.bzl b/python/uv/private/lock.bzl
@@ -16,10 +16,8 @@
 """
 
 load("@bazel_skylib//lib:shell.bzl", "shell")
-load("@pythons_hub//:versions.bzl", "DEFAULT_PYTHON_VERSION", "MINOR_MAPPING")
 load("//python:py_binary.bzl", "py_binary")
 load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
-load("//python/private:full_version.bzl", "full_version")
 load("//python/private:toolchain_types.bzl", "EXEC_TOOLS_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
 load(":toolchain_types.bzl", "UV_TOOLCHAIN_TYPE")
 
@@ -75,15 +73,15 @@ def _args(ctx):
 
 def _lock_impl(ctx):
     srcs = ctx.files.srcs
-    python_version = full_version(
-        version = ctx.attr.python_version or DEFAULT_PYTHON_VERSION,
-        minor_mapping = MINOR_MAPPING,
-    )
-    output = ctx.actions.declare_file("{}.{}.out".format(
-        ctx.label.name,
-        python_version.replace(".", "_"),
-    ))
+    fname = "{}.out".format(ctx.label.name)
+    python_version = ctx.attr.python_version
+    if python_version:
+        fname = "{}.{}.out".format(
+            ctx.label.name,
+            python_version.replace(".", "_"),
+        )
 
+    output = ctx.actions.declare_file(fname)
     toolchain_info = ctx.toolchains[UV_TOOLCHAIN_TYPE]
     uv = toolchain_info.uv_toolchain_info.uv[DefaultInfo].files_to_run.executable
 
@@ -166,15 +164,7 @@ def _transition_impl(input_settings, attr):
         _PYTHON_VERSION_FLAG: input_settings[_PYTHON_VERSION_FLAG],
     }
     if attr.python_version:
-        # FIXME @aignas 2025-03-20: using `full_version` is a workaround for a bug in
-        # how we order toolchains in bazel. If I set the `python_version` flag
-        # to `3.12`, I would expect the latest version to be selected, i.e. the
-        # one that is in MINOR_MAPPING, but it seems that 3.12.0 is selected,
-        # because of how the targets are ordered.
-        settings[_PYTHON_VERSION_FLAG] = full_version(
-            version = attr.python_version,
-            minor_mapping = MINOR_MAPPING,
-        )
+        settings[_PYTHON_VERSION_FLAG] = attr.python_version
     return settings
 
 _python_version_transition = transition(
@@ -436,9 +426,6 @@ def lock(
     if not BZLMOD_ENABLED:
         kwargs["target_compatible_with"] = ["@platforms//:incompatible"]
 
-    # FIXME @aignas 2025-03-17: should we have one more target that transitions
-    # the python_version to ensure that if somebody calls `bazel build
-    # :requirements` that it is locked with the right `python_version`?
     _lock(
         name = name,
         args = args,
diff --git a/tests/config_settings/transition/multi_version_tests.bzl b/tests/config_settings/transition/multi_version_tests.bzl
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for py_test."""
 
+load("@pythons_hub//:versions.bzl", "DEFAULT_PYTHON_VERSION")
 load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("@rules_testing//lib:util.bzl", "TestingAspectInfo", rt_util = "util")
@@ -29,7 +30,7 @@ load("//tests/support:support.bzl", "CC_TOOLCHAIN")
 # If the toolchain is not resolved then you will have a weird message telling
 # you that your transition target does not have a PyRuntime provider, which is
 # caused by there not being a toolchain detected for the target.
-_PYTHON_VERSION = "3.11"
+_PYTHON_VERSION = DEFAULT_PYTHON_VERSION
 
 _tests = []
 
diff --git a/tests/python/python_tests.bzl b/tests/python/python_tests.bzl
@@ -284,6 +284,58 @@ def _test_default_non_rules_python_ignore_root_user_error_non_root_module(env):
 
 _tests.append(_test_default_non_rules_python_ignore_root_user_error_non_root_module)
 
+def _test_toolchain_ordering(env):
+    py = parse_modules(
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "my_module",
+                toolchain = [
+                    _toolchain("3.10"),
+                    _toolchain("3.10.15"),
+                    _toolchain("3.10.16"),
+                    _toolchain("3.10.11"),
+                    _toolchain("3.11.1"),
+                    _toolchain("3.11.10"),
+                    _toolchain("3.11.11", is_default = True),
+                ],
+            ),
+            _mod(name = "rules_python", toolchain = [_toolchain("3.11")]),
+        ),
+    )
+    got_versions = [
+        t.python_version
+        for t in py.toolchains
+    ]
+
+    env.expect.that_str(py.default_python_version).equals("3.11.11")
+    env.expect.that_dict(py.config.minor_mapping).contains_exactly({
+        "3.10": "3.10.16",
+        "3.11": "3.11.11",
+        "3.12": "3.12.9",
+        "3.13": "3.13.2",
+        "3.8": "3.8.20",
+        "3.9": "3.9.21",
+    })
+    env.expect.that_collection(got_versions).contains_exactly([
+        # First the full-version toolchains that are in minor_mapping
+        # so that they get matched first if only the `python_version` is in MINOR_MAPPING
+        #
+        # The default version is always set in the `python_version` flag, so know, that
+        # the default match will be somewhere in the first bunch.
+        "3.10",
+        "3.10.16",
+        "3.11",
+        "3.11.11",
+        # Next, the rest, where we will match things based on the `python_version` being
+        # the same
+        "3.10.15",
+        "3.10.11",
+        "3.11.1",
+        "3.11.10",
+    ]).in_order()
+
+_tests.append(_test_toolchain_ordering)
+
 def _test_default_from_defaults(env):
     py = parse_modules(
         module_ctx = _mock_mctx(
diff --git a/tests/toolchains/transitions/BUILD.bazel b/tests/toolchains/transitions/BUILD.bazel
@@ -0,0 +1,5 @@
+load(":transitions_tests.bzl", "transitions_test_suite")
+
+transitions_test_suite(
+    name = "transitions_tests",
+)
diff --git a/tests/toolchains/transitions/transitions_tests.bzl b/tests/toolchains/transitions/transitions_tests.bzl
@@ -0,0 +1,182 @@
+# Copyright 2022 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.
+
+""
+
+load("@pythons_hub//:versions.bzl", "DEFAULT_PYTHON_VERSION", "MINOR_MAPPING")
+load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("@rules_testing//lib:util.bzl", rt_util = "util")
+load("//python:versions.bzl", "TOOL_VERSIONS")
+load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
+load("//python/private:full_version.bzl", "full_version")  # buildifier: disable=bzl-visibility
+load("//python/private:toolchain_types.bzl", "EXEC_TOOLS_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
+load("//tests/support:support.bzl", "PYTHON_VERSION")
+
+_analysis_tests = []
+
+def _transition_impl(input_settings, attr):
+    """Transition based on python_version flag.
+
+    This is a simple transition impl that a user of rules_python may implement
+    for their own rule.
+    """
+    settings = {
+        PYTHON_VERSION: input_settings[PYTHON_VERSION],
+    }
+    if attr.python_version:
+        settings[PYTHON_VERSION] = attr.python_version
+    return settings
+
+_python_version_transition = transition(
+    implementation = _transition_impl,
+    inputs = [PYTHON_VERSION],
+    outputs = [PYTHON_VERSION],
+)
+
+TestInfo = provider(
+    doc = "A simple test provider to forward the values for the assertion.",
+    fields = {"got": "", "want": ""},
+)
+
+def _impl(ctx):
+    if ctx.attr.skip:
+        return [TestInfo(got = "", want = "")]
+
+    exec_tools = ctx.toolchains[EXEC_TOOLS_TOOLCHAIN_TYPE].exec_tools
+    got_version = exec_tools.exec_interpreter[platform_common.ToolchainInfo].py3_runtime.interpreter_version_info
+
+    return [
+        TestInfo(
+            got = "{}.{}.{}".format(
+                got_version.major,
+                got_version.minor,
+                got_version.micro,
+            ),
+            want = ctx.attr.want_version,
+        ),
+    ]
+
+_simple_transition = rule(
+    implementation = _impl,
+    attrs = {
+        "python_version": attr.string(
+            doc = "The input python version which we transition on.",
+        ),
+        "skip": attr.bool(
+            doc = "Whether to skip the test",
+        ),
+        "want_version": attr.string(
+            doc = "The python version that we actually expect to receive.",
+        ),
+        "_allowlist_function_transition": attr.label(
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+        ),
+    },
+    toolchains = [
+        config_common.toolchain_type(
+            EXEC_TOOLS_TOOLCHAIN_TYPE,
+            mandatory = False,
+        ),
+    ],
+    cfg = _python_version_transition,
+)
+
+def _test_transitions(*, name, tests, skip = False):
+    """A reusable rule so that we can split the tests."""
+    targets = {}
+    for test_name, (input_version, want_version) in tests.items():
+        target_name = "{}_{}".format(name, test_name)
+        targets["python_" + test_name] = target_name
+        rt_util.helper_target(
+            _simple_transition,
+            name = target_name,
+            python_version = input_version,
+            want_version = want_version,
+            skip = skip,
+        )
+
+    analysis_test(
+        name = name,
+        impl = _test_transition_impl,
+        targets = targets,
+    )
+
+def _test_transition_impl(env, targets):
+    # Check that the forwarded version from the PyRuntimeInfo is correct
+    for target in dir(targets):
+        if not target.startswith("python"):
+            # Skip other attributes that might be not the ones we set (e.g. to_json, to_proto).
+            continue
+
+        test_info = env.expect.that_target(getattr(targets, target)).provider(
+            TestInfo,
+            factory = lambda v, meta: v,
+        )
+        env.expect.that_str(test_info.got).equals(test_info.want)
+
+def _test_full_version(name):
+    """Check that python_version transitions work.
+
+    Expectation is to get the same full version that we input.
+    """
+    _test_transitions(
+        name = name,
+        tests = {
+            v.replace(".", "_"): (v, v)
+            for v in TOOL_VERSIONS
+        },
+    )
+
+_analysis_tests.append(_test_full_version)
+
+def _test_minor_versions(name):
+    """Ensure that MINOR_MAPPING versions are correctly selected."""
+    _test_transitions(
+        name = name,
+        skip = not BZLMOD_ENABLED,
+        tests = {
+            minor.replace(".", "_"): (minor, full)
+            for minor, full in MINOR_MAPPING.items()
+        },
+    )
+
+_analysis_tests.append(_test_minor_versions)
+
+def _test_default(name):
+    """Check the default version.
+
+    Lastly, if we don't provide any version to the transition, we should
+    get the default version
+    """
+    default_version = full_version(
+        version = DEFAULT_PYTHON_VERSION,
+        minor_mapping = MINOR_MAPPING,
+    ) if DEFAULT_PYTHON_VERSION else ""
+
+    _test_transitions(
+        name = name,
+        skip = not BZLMOD_ENABLED,
+        tests = {
+            "default": (None, default_version),
+        },
+    )
+
+_analysis_tests.append(_test_default)
+
+def transitions_test_suite(name):
+    test_suite(
+        name = name,
+        tests = _analysis_tests,
+    )
