feat(bzlmod): add python.override APIs

Before this PR the users could not override how the hermetic toolchain
is downloaded when in `bzlmod`. However, the same APIs would be
available to users using `WORKSPACE`. With this we also allow root
modules to restrict which toolchain versions the non-root modules, which
may be helpful when optimizing the CI runtimes so that we don't waste
time downloading multiple `micro` versions of the same `3.X` python
version, which most of the times have identical behavior.

Work towards bazelbuild#2081 and this should be one of the last items that are
blocking bazelbuild#1361 from the API point of view.

.bazelrc

@@ -4,8 +4,8 @@
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, execute
 # `bazel run @rules_bazel_integration_test//tools:update_deleted_packages`
-build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/python/private,gazelle/pythonconfig,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/python/private,gazelle/pythonconfig,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered
 
 test --test_output=errors
 

CHANGELOG.md

@@ -28,6 +28,10 @@ A brief description of the categories of changes:
 * (gazelle): Update error messages when unable to resolve a dependency to be more human-friendly.
 * (flags) The {obj}`--python_version` flag now also returns
   {obj}`config_common.FeatureFlagInfo`.
+* (toolchain): The toolchain patches now expose the `patch_strip` attribute
+  that one should use when patching toolchains. Please set it if you are
+  patching python interpreter. In the next release the default will be set to
+  `0` which better reflects the defaults used in public `bazel` APIs.
 * (toolchains) When {obj}`py_runtime.interpreter_version_info` isn't specified,
   the {obj}`--python_version` flag will determine the value. This allows
   specifying the build-time Python version for the
@@ -60,6 +64,10 @@ A brief description of the categories of changes:
 
 
 ### Added
+* (bzlmod): Toolchain overrides can now be done using the new
+  {bzl:obj}`python.override`, {bzl:obj}`python.single_version_override` and
+  {bzl:obj}`python.single_version_platform_override` tag classes.
+  See [#2081](https://github.com/bazelbuild/rules_python/issues/2081).
 * (rules) Executables provide {obj}`PyExecutableInfo`, which contains
   executable-specific information useful for packaging an executable or
   or deriving a new one from the original.

MODULE.bazel

@@ -81,7 +81,7 @@ dev_python = use_extension(
     "python",
     dev_dependency = True,
 )
-dev_python.rules_python_private_testing(
+dev_python.override(
     register_all_versions = True,
 )
 

docs/toolchains.md

@@ -164,6 +164,17 @@ Remember to call `use_repo()` to make repos visible to your module:
 Python toolchains can be utilized in other bazel rules, such as `genrule()`, by adding the `toolchains=["@rules_python//python:current_py_toolchain"]` attribute. You can obtain the path to the Python interpreter using the `$(PYTHON2)` and `$(PYTHON3)` ["Make" Variables](https://bazel.build/reference/be/make-variables). See the
 {gh-path}`test_current_py_toolchain <tests/load_from_macro/BUILD.bazel>` target for an example.
 
+### Overriding toolchain defaults and adding more toolchains
+
+One can perform various overrides for the registered toolchains from the root module. For example, the following usecases would be supported using the existing attributes:
+
+* Limiting the available toolchains for the entire `bzlmod` transitive graph
+  via {attr}`python.override.available_python_versions`.
+* Setting particular `X.Y.Z` python versions when modules request `X.Y` version
+  via {attr}`python.override.minor_mapping`.
+* Adding custom {attr}`python.single_version_platform_override.coverage_tool`.
+* Adding new python versions via {bzl:obj}`python.single_version_override` or
+  {bzl:obj}`python.single_version_platform_override`.
 
 ## Workspace configuration
 
@@ -240,5 +251,5 @@ automatically registers a higher-priority toolchain; it won't be used unless
 there is a toolchain misconfiguration somewhere.
 
 To aid migration off the Bazel-builtin toolchain, rules_python provides
-{obj}`@rules_python//python/runtime_env_toolchains:all`. This is an equivalent
+{bzl:obj}`@rules_python//python/runtime_env_toolchains:all`. This is an equivalent
 toolchain, but is implemented using rules_python's objects.

examples/bzlmod/MODULE.bazel

@@ -22,7 +22,7 @@ bazel_dep(name = "protobuf", version = "24.4", repo_name = "com_google_protobuf"
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
     configure_coverage_tool = True,
-    # Only set when you have mulitple toolchain versions.
+    # Only set when you have multiple toolchain versions.
     is_default = True,
     python_version = "3.9",
 )
@@ -37,6 +37,54 @@ python.toolchain(
     python_version = "3.10",
 )
 
+# One can override the actual toolchain versions that are available, which can be useful
+# to when optimizing what gets downloaded and when.
+python.override(
+    available_python_versions = [
+        "3.10.9",
+        "3.9.19",
+        # The following is used by the `other_module` and we need to include it here
+        # as well.
+        "3.11.8",
+    ],
+    # Also override the `minor_mapping` so that when the modules specify a particular
+    # `3.X` version, we decide what gets used.
+    minor_mapping = {
+        "3.10": "3.10.9",
+        "3.11": "3.11.8",
+        "3.9": "3.9.19",
+    },
+)
+
+# Or the sources that the toolchains come from for all platforms
+python.single_version_override(
+    patch_strip = 1,
+    # The user can specify patches to be applied to all interpreters.
+    patches = [],
+    python_version = "3.10.2",
+    sha256 = {
+        "aarch64-apple-darwin": "1409acd9a506e2d1d3b65c1488db4e40d8f19d09a7df099667c87a506f71c0ef",
+        "aarch64-unknown-linux-gnu": "8f351a8cc348bb45c0f95b8634c8345ec6e749e483384188ad865b7428342703",
+        "x86_64-apple-darwin": "8146ad4390710ec69b316a5649912df0247d35f4a42e2aa9615bffd87b3e235a",
+        "x86_64-pc-windows-msvc": "a1d9a594cd3103baa24937ad9150c1a389544b4350e859200b3e5c036ac352bd",
+        "x86_64-unknown-linux-gnu": "9b64eca2a94f7aff9409ad70bdaa7fbbf8148692662e764401883957943620dd",
+    },
+    urls = ["20220227/cpython-{python_version}+20220227-{platform}-{build}.tar.gz"],
+)
+
+# Or a single platform. This can be used in combination with the
+# `single_version_override` and `single_version_platform_override` will be
+# applied after `single_version_override`. Any values present in this override
+# will overwrite the values set by the `single_version_override`
+python.single_version_platform_override(
+    patch_strip = 1,
+    patches = [],
+    platform = "aarch64-apple-darwin",
+    python_version = "3.10.2",
+    sha256 = "1409acd9a506e2d1d3b65c1488db4e40d8f19d09a7df099667c87a506f71c0ef",
+    urls = ["20220227/cpython-{python_version}+20220227-{platform}-{build}.tar.gz"],
+)
+
 # You only need to load this repositories if you are using multiple Python versions.
 # See the tests folder for various examples on using multiple Python versions.
 # The names "python_3_9" and "python_3_10" are autmatically created by the repo

python/extensions/python.bzl

@@ -12,7 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"Python toolchain module extensions for use with bzlmod"
+"""Python toolchain module extensions for use with bzlmod.
+
+## Basic usage
+
+The simplest way to configure the toolchain with `rules_python` is as follows.
+
+```starlark
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(
+    is_default = True,
+    python_version = "3.11",
+)
+use_repo(python, "python_3_11")
+```
+
+For more in-depth documentation see the {rule}`python.toolchain`.
+
+## Overrides
+
+Overrides can be done at 3 different levels:
+* Overrides affecting all python toolchain versions on all platforms - {obj}`python.override`.
+* Overrides affecting a single toolchain versions on all platforms - {obj}`python.single_version_override`.
+* Overrides affecting a single toolchain versions on a single platforms - {obj}`python.single_version_platform_override`.
+
+:::{seealso}
+The main documentation page on registering [toolchains](/toolchains).
+:::
+"""
 
 load("//python/private:python.bzl", _python = "python")
 

python/private/BUILD.bazel

@@ -132,11 +132,12 @@ bzl_library(
     srcs = ["python.bzl"],
     deps = [
         ":full_version_bzl",
+        ":python_repositories_bzl",
         ":pythons_hub_bzl",
         ":repo_utils_bzl",
+        ":semver_bzl",
         ":toolchains_repo_bzl",
         ":util_bzl",
-        "//python:repositories_bzl",
         "@bazel_features//:features",
     ],
 )
@@ -162,7 +163,6 @@ bzl_library(
     name = "pythons_hub_bzl",
     srcs = ["pythons_hub.bzl"],
     deps = [
-        ":full_version_bzl",
         ":py_toolchain_suite_bzl",
     ],
 )

python/private/common/py_runtime_rule.bzl

@@ -202,8 +202,8 @@ See @bazel_tools//tools/python:python_bootstrap_template.txt for more variables.
         "coverage_tool": attr.label(
             allow_files = False,
             doc = """
-This is a target to use for collecting code coverage information from `py_binary`
-and `py_test` targets.
+This is a target to use for collecting code coverage information from
+{rule}`py_binary` and {rule}`py_test` targets.
 
 If set, the target must either produce a single file or be an executable target.
 The path to the single file, or the executable if the target is executable,
@@ -212,7 +212,7 @@ runfiles will be added to the runfiles when coverage is enabled.
 
 The entry point for the tool must be loadable by a Python interpreter (e.g. a
 `.py` or `.pyc` file).  It must accept the command line arguments
-of coverage.py (https://coverage.readthedocs.io), at least including
+of [`coverage.py`](https://coverage.readthedocs.io), at least including
 the `run` and `lcov` subcommands.
 """,
         ),
@@ -311,7 +311,7 @@ The template to use when two stage bootstrapping is enabled
             default = DEFAULT_STUB_SHEBANG,
             doc = """
 "Shebang" expression prepended to the bootstrapping Python stub script
-used when executing `py_binary` targets.
+used when executing {rule}`py_binary` targets.
 
 See https://github.com/bazelbuild/bazel/issues/8685 for
 motivation.

python/private/pypi/BUILD.bazel

@@ -59,7 +59,6 @@ bzl_library(
     srcs = ["extension.bzl"],
     deps = [
         ":attrs_bzl",
-        "//python/private:semver_bzl",
         ":hub_repository_bzl",
         ":parse_requirements_bzl",
         ":evaluate_markers_bzl",
@@ -71,6 +70,7 @@ bzl_library(
         "//python/private:full_version_bzl",
         "//python/private:normalize_name_bzl",
         "//python/private:version_label_bzl",
+        "//python/private:semver_bzl",
         "@bazel_features//:features",
     ] + [
         "@pythons_hub//:interpreters_bzl",