feat(python.toolchain): support file-based default Python version (bazel-contrib#2588)

This change adds a new `default_version_file` attribute to
`python.toolchain`. If set, the toolchain compares the file's contents
to its `python_version`, and if they match, treats that toolchain as
default (ignoring `is_default`). This allows Bazel to synchronize the
default Python version with external tools (e.g., pyenv) that use a
`.python-version` file or environment variables.

Fixes bazel-contrib#2587.

---------

Co-authored-by: Ignas Anikevicius <[email protected]>

Author: vonschultz

CHANGELOG.md

@@ -128,6 +128,9 @@ Unreleased changes template.
 
 {#v1-3-0-added}
 ### Added
+* (python) {attr}`python.defaults` has been added to allow users to
+  set the default python version in the root module by reading the
+  default version number from a file or an environment variable.
 * {obj}`//python/bin:python`: convenience target for directly running an
   interpreter. {obj}`--//python/bin:python_src` can be used to specify a
   binary whose interpreter to use.

examples/multi_python_versions/MODULE.bazel

@@ -10,6 +10,11 @@ local_path_override(
 )
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.defaults(
+    # The environment variable takes precedence if set.
+    python_version = "3.9",
+    python_version_env = "BAZEL_PYTHON_VERSION",
+)
 python.toolchain(
     configure_coverage_tool = True,
     # Only set when you have mulitple toolchain versions.

python/private/python.bzl

@@ -78,6 +78,47 @@ def parse_modules(*, module_ctx, _fail = fail):
 
     config = _get_toolchain_config(modules = module_ctx.modules, _fail = _fail)
 
+    default_python_version = None
+    for mod in module_ctx.modules:
+        defaults_attr_structs = _create_defaults_attr_structs(mod = mod)
+        default_python_version_env = None
+        default_python_version_file = None
+
+        # Only the root module and rules_python are allowed to specify the default
+        # toolchain for a couple reasons:
+        # * It prevents submodules from specifying different defaults and only
+        #   one of them winning.
+        # * rules_python needs to set a soft default in case the root module doesn't,
+        #   e.g. if the root module doesn't use Python itself.
+        # * The root module is allowed to override the rules_python default.
+        if mod.is_root or (mod.name == "rules_python" and not default_python_version):
+            for defaults_attr in defaults_attr_structs:
+                default_python_version = _one_or_the_same(
+                    default_python_version,
+                    defaults_attr.python_version,
+                    onerror = _fail_multiple_defaults_python_version,
+                )
+                default_python_version_env = _one_or_the_same(
+                    default_python_version_env,
+                    defaults_attr.python_version_env,
+                    onerror = _fail_multiple_defaults_python_version_env,
+                )
+                default_python_version_file = _one_or_the_same(
+                    default_python_version_file,
+                    defaults_attr.python_version_file,
+                    onerror = _fail_multiple_defaults_python_version_file,
+                )
+            if default_python_version_file:
+                default_python_version = _one_or_the_same(
+                    default_python_version,
+                    module_ctx.read(default_python_version_file, watch = "yes").strip(),
+                )
+            if default_python_version_env:
+                default_python_version = module_ctx.getenv(
+                    default_python_version_env,
+                    default_python_version,
+                )
+
     seen_versions = {}
     for mod in module_ctx.modules:
         module_toolchain_versions = []
@@ -104,7 +145,13 @@ def parse_modules(*, module_ctx, _fail = fail):
                 # * rules_python needs to set a soft default in case the root module doesn't,
                 #   e.g. if the root module doesn't use Python itself.
                 # * The root module is allowed to override the rules_python default.
-                is_default = toolchain_attr.is_default
+                if default_python_version:
+                    is_default = default_python_version == toolchain_version
+                    if toolchain_attr.is_default and not is_default:
+                        fail("The 'is_default' attribute doesn't work if you set " +
+                             "the default Python version with the `defaults` tag.")
+                else:
+                    is_default = toolchain_attr.is_default
 
                 # Also only the root module should be able to decide ignore_root_user_error.
                 # Modules being depended upon don't know the final environment, so they aren't
@@ -115,7 +162,7 @@ def parse_modules(*, module_ctx, _fail = fail):
                     fail("Toolchains in the root module must have consistent 'ignore_root_user_error' attributes")
 
                 ignore_root_user_error = toolchain_attr.ignore_root_user_error
-            elif mod.name == "rules_python" and not default_toolchain:
+            elif mod.name == "rules_python" and not default_toolchain and not default_python_version:
                 # We don't do the len() check because we want the default that rules_python
                 # sets to be clearly visible.
                 is_default = toolchain_attr.is_default
@@ -282,6 +329,19 @@ def _python_impl(module_ctx):
     else:
         return None
 
+def _one_or_the_same(first, second, *, onerror = None):
+    if not first:
+        return second
+    if not second or second == first:
+        return first
+    if onerror:
+        return onerror(first, second)
+    else:
+        fail("Unique value needed, got both '{}' and '{}', which are different".format(
+            first,
+            second,
+        ))
+
 def _fail_duplicate_module_toolchain_version(version, module):
     fail(("Duplicate module toolchain version: module '{module}' attempted " +
           "to use version '{version}' multiple times in itself").format(
@@ -305,6 +365,30 @@ def _warn_duplicate_global_toolchain_version(version, first, second_toolchain_na
         version = version,
     ))
 
+def _fail_multiple_defaults_python_version(first, second):
+    fail(("Multiple python_version entries in defaults: " +
+          "First default was python_version '{first}'. " +
+          "Second was python_version '{second}'").format(
+        first = first,
+        second = second,
+    ))
+
+def _fail_multiple_defaults_python_version_file(first, second):
+    fail(("Multiple python_version_file entries in defaults: " +
+          "First default was python_version_file '{first}'. " +
+          "Second was python_version_file '{second}'").format(
+        first = first,
+        second = second,
+    ))
+
+def _fail_multiple_defaults_python_version_env(first, second):
+    fail(("Multiple python_version_env entries in defaults: " +
+          "First default was python_version_env '{first}'. " +
+          "Second was python_version_env '{second}'").format(
+        first = first,
+        second = second,
+    ))
+
 def _fail_multiple_default_toolchains(first, second):
     fail(("Multiple default toolchains: only one toolchain " +
           "can have is_default=True. First default " +
@@ -526,6 +610,21 @@ def _get_toolchain_config(*, modules, _fail = fail):
         register_all_versions = register_all_versions,
     )
 
+def _create_defaults_attr_structs(*, mod):
+    arg_structs = []
+
+    for tag in mod.tags.defaults:
+        arg_structs.append(_create_defaults_attr_struct(tag = tag))
+
+    return arg_structs
+
+def _create_defaults_attr_struct(*, tag):
+    return struct(
+        python_version = getattr(tag, "python_version", None),
+        python_version_env = getattr(tag, "python_version_env", None),
+        python_version_file = getattr(tag, "python_version_file", None),
+    )
+
 def _create_toolchain_attr_structs(*, mod, config, seen_versions):
     arg_structs = []
 
@@ -570,6 +669,49 @@ def _get_bazel_version_specific_kwargs():
 
     return kwargs
 
+_defaults = tag_class(
+    doc = """Tag class to specify the default Python version.""",
+    attrs = {
+        "python_version": attr.string(
+            mandatory = False,
+            doc = """\
+String saying what the default Python version should be. If the string
+matches the {attr}`python_version` attribute of a toolchain, this
+toolchain is the default version. If this attribute is set, the
+{attr}`is_default` attribute of the toolchain is ignored.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+        ),
+        "python_version_env": attr.string(
+            mandatory = False,
+            doc = """\
+Environment variable saying what the default Python version should be.
+If the string matches the {attr}`python_version` attribute of a
+toolchain, this toolchain is the default version. If this attribute is
+set, the {attr}`is_default` attribute of the toolchain is ignored.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+        ),
+        "python_version_file": attr.label(
+            mandatory = False,
+            allow_single_file = True,
+            doc = """\
+File saying what the default Python version should be. If the contents
+of the file match the {attr}`python_version` attribute of a toolchain,
+this toolchain is the default version. If this attribute is set, the
+{attr}`is_default` attribute of the toolchain is ignored.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+        ),
+    },
+)
+
 _toolchain = tag_class(
     doc = """Tag class used to register Python toolchains.
 Use this tag class to register one or more Python toolchains. This class
@@ -653,7 +795,14 @@ error to run with root access instead.
         ),
         "is_default": attr.bool(
             mandatory = False,
-            doc = "Whether the toolchain is the default version",
+            doc = """\
+Whether the toolchain is the default version.
+
+:::{versionchanged} VERSION_NEXT_FEATURE
+This setting is ignored if the default version is set using the `defaults`
+tag class.
+:::
+""",
         ),
         "python_version": attr.string(
             mandatory = True,
@@ -852,6 +1001,7 @@ python = module_extension(
 """,
     implementation = _python_impl,
     tag_classes = {
+        "defaults": _defaults,
         "override": _override,
         "single_version_override": _single_version_override,
         "single_version_platform_override": _single_version_platform_override,

tests/python/python_tests.bzl

@@ -20,8 +20,11 @@ load("//python/private:python.bzl", "parse_modules")  # buildifier: disable=bzl-
 
 _tests = []
 
-def _mock_mctx(*modules, environ = {}):
+def _mock_mctx(*modules, environ = {}, mocked_files = {}):
     return struct(
+        path = lambda x: struct(exists = x in mocked_files, _file = x),
+        read = lambda x, watch = None: mocked_files[x._file if "_file" in dir(x) else x],
+        getenv = environ.get,
         os = struct(environ = environ),
         modules = [
             struct(
@@ -39,10 +42,11 @@ def _mock_mctx(*modules, environ = {}):
         ],
     )
 
-def _mod(*, name, toolchain = [], override = [], single_version_override = [], single_version_platform_override = [], is_root = True):
+def _mod(*, name, defaults = [], toolchain = [], override = [], single_version_override = [], single_version_platform_override = [], is_root = True):
     return struct(
         name = name,
         tags = struct(
+            defaults = defaults,
             toolchain = toolchain,
             override = override,
             single_version_override = single_version_override,
@@ -51,6 +55,13 @@ def _mod(*, name, toolchain = [], override = [], single_version_override = [], s
         is_root = is_root,
     )
 
+def _defaults(python_version = None, python_version_env = None, python_version_file = None):
+    return struct(
+        python_version = python_version,
+        python_version_env = python_version_env,
+        python_version_file = python_version_file,
+    )
+
 def _toolchain(python_version, *, is_default = False, **kwargs):
     return struct(
         is_default = is_default,
@@ -273,6 +284,86 @@ 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_default_from_defaults(env):
+    py = parse_modules(
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "my_root_module",
+                defaults = [_defaults(python_version = "3.11")],
+                toolchain = [_toolchain("3.10"), _toolchain("3.11"), _toolchain("3.12")],
+                is_root = True,
+            ),
+        ),
+    )
+
+    env.expect.that_str(py.default_python_version).equals("3.11")
+
+    want_toolchains = [
+        struct(
+            name = "python_3_" + minor_version,
+            python_version = "3." + minor_version,
+            register_coverage_tool = False,
+        )
+        for minor_version in ["10", "11", "12"]
+    ]
+    env.expect.that_collection(py.toolchains).contains_exactly(want_toolchains)
+
+_tests.append(_test_default_from_defaults)
+
+def _test_default_from_defaults_env(env):
+    py = parse_modules(
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "my_root_module",
+                defaults = [_defaults(python_version = "3.11", python_version_env = "PYENV_VERSION")],
+                toolchain = [_toolchain("3.10"), _toolchain("3.11"), _toolchain("3.12")],
+                is_root = True,
+            ),
+            environ = {"PYENV_VERSION": "3.12"},
+        ),
+    )
+
+    env.expect.that_str(py.default_python_version).equals("3.12")
+
+    want_toolchains = [
+        struct(
+            name = "python_3_" + minor_version,
+            python_version = "3." + minor_version,
+            register_coverage_tool = False,
+        )
+        for minor_version in ["10", "11", "12"]
+    ]
+    env.expect.that_collection(py.toolchains).contains_exactly(want_toolchains)
+
+_tests.append(_test_default_from_defaults_env)
+
+def _test_default_from_defaults_file(env):
+    py = parse_modules(
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "my_root_module",
+                defaults = [_defaults(python_version_file = "@@//:.python-version")],
+                toolchain = [_toolchain("3.10"), _toolchain("3.11"), _toolchain("3.12")],
+                is_root = True,
+            ),
+            mocked_files = {"@@//:.python-version": "3.12\n"},
+        ),
+    )
+
+    env.expect.that_str(py.default_python_version).equals("3.12")
+
+    want_toolchains = [
+        struct(
+            name = "python_3_" + minor_version,
+            python_version = "3." + minor_version,
+            register_coverage_tool = False,
+        )
+        for minor_version in ["10", "11", "12"]
+    ]
+    env.expect.that_collection(py.toolchains).contains_exactly(want_toolchains)
+
+_tests.append(_test_default_from_defaults_file)
+
 def _test_first_occurance_of_the_toolchain_wins(env):
     py = parse_modules(
         module_ctx = _mock_mctx(