refactor(zipapp): add dev-only toolchain for the Rust exe_zip_maker

MODULE.bazel

@@ -112,6 +112,12 @@ rust_crates.spec(
 rust_crates.from_specs()
 use_repo(rust_crates, "crates")
 
+# Use the Rust-based build tools when developing rules_python itself.
+register_toolchains(
+    "//dev/dev_only_toolchains:exe_zip_maker_toolchain",
+    dev_dependency = True,
+)
+
 bazel_dep(name = "other", version = "0", dev_dependency = True)
 bazel_dep(name = "another_module", version = "0", dev_dependency = True)
 

dev/dev_only_toolchains/BUILD.bazel

@@ -0,0 +1,39 @@
+# Toolchains that are only registered when developing rules_python itself
+# (see `register_toolchains(..., dev_dependency = True)` in MODULE.bazel).
+# They are not registered in WORKSPACE mode.
+#
+# Only `toolchain()` declarations belong here; implementations live in
+# `impls/` so that registering the toolchains doesn't load the implementations.
+
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+
+package(default_visibility = ["//visibility:private"])
+
+# Whether to use the Rust implementation of exe_zip_maker.
+# * `yes`: use it.
+# * `no`: use the built-in Python implementation.
+# * `auto`: let rules_python decide. Currently the same as `no`.
+#
+# Enable with `--//dev/dev_only_toolchains:use_rust_exe_zip_maker=yes`.
+string_flag(
+    name = "use_rust_exe_zip_maker",
+    build_setting_default = "auto",
+    values = [
+        "auto",
+        "no",
+        "yes",
+    ],
+    visibility = ["//tests:__subpackages__"],
+)
+
+config_setting(
+    name = "is_rust_exe_zip_maker_yes",
+    flag_values = {":use_rust_exe_zip_maker": "yes"},
+)
+
+toolchain(
+    name = "exe_zip_maker_toolchain",
+    target_settings = [":is_rust_exe_zip_maker_yes"],
+    toolchain = "//dev/dev_only_toolchains/impls:exe_zip_maker",
+    toolchain_type = "//python/private/toolchain_types:exe_zip_maker",
+)

dev/dev_only_toolchains/impls/BUILD.bazel

@@ -0,0 +1,15 @@
+# Toolchain implementations for the dev-only toolchains.
+#
+# These are kept in a separate package from the `toolchain()` declarations so
+# that toolchain resolution only needs to load this package (and the tool's
+# dependencies) when a toolchain is actually selected.
+
+load("//python/private/zipapp:py_exe_zip_maker_toolchain.bzl", "py_exe_zip_maker_toolchain")  # buildifier: disable=bzl-visibility
+
+package(default_visibility = ["//dev/dev_only_toolchains:__pkg__"])
+
+# Use the Rust implementation of exe_zip_maker instead of the Python one.
+py_exe_zip_maker_toolchain(
+    name = "exe_zip_maker",
+    exe_zip_maker = "//crates/exe_zip_maker",
+)

python/private/py_executable.bzl

@@ -68,7 +68,13 @@ load(":py_internal.bzl", "py_internal")
 load(":py_runtime_info.bzl", "DEFAULT_STUB_SHEBANG")
 load(":reexports.bzl", "BuiltinPyInfo", "BuiltinPyRuntimeInfo")
 load(":rule_builders.bzl", "ruleb")
-load(":toolchain_types.bzl", "CC_TOOLCHAIN_TYPE", "EXEC_TOOLS_TOOLCHAIN_TYPE", "LAUNCHER_MAKER_TOOLCHAIN_TYPE", TOOLCHAIN_TYPE = "TARGET_TOOLCHAIN_TYPE")
+load(
+    ":toolchain_types.bzl",
+    "CC_TOOLCHAIN_TYPE",
+    "EXEC_TOOLS_TOOLCHAIN_TYPE",
+    "LAUNCHER_MAKER_TOOLCHAIN_TYPE",
+    TOOLCHAIN_TYPE = "TARGET_TOOLCHAIN_TYPE",
+)
 load(":transition_labels.bzl", "TRANSITION_LABELS")
 load(":venv_runfiles.bzl", "create_venv_app_files")
 

python/private/toolchain_types.bzl

@@ -20,6 +20,7 @@ implementation of the toolchain.
 
 TARGET_TOOLCHAIN_TYPE = Label("//python:toolchain_type")
 EXEC_TOOLS_TOOLCHAIN_TYPE = Label("//python:exec_tools_toolchain_type")
+EXE_ZIP_MAKER_TOOLCHAIN_TYPE = Label("//python/private/toolchain_types:exe_zip_maker")
 PY_CC_TOOLCHAIN_TYPE = Label("//python/cc:toolchain_type")
 LAUNCHER_MAKER_TOOLCHAIN_TYPE = Label("@bazel_tools//tools/launcher:launcher_maker_toolchain_type")
 CC_TOOLCHAIN_TYPE = Label("@bazel_tools//tools/cpp:toolchain_type")

python/private/toolchain_types/BUILD.bazel

@@ -0,0 +1,22 @@
+# Toolchain types used internally by rules_python.
+#
+# These are in a separate package so toolchain() registrations and rules can
+# depend on the type without loading the toolchain implementation.
+
+load("//python/private:distribution_filegroup.bzl", "distribution_filegroup")
+
+package(default_visibility = ["//:__subpackages__"])
+
+licenses(["notice"])
+
+distribution_filegroup(
+    name = "distribution",
+)
+
+# Toolchain providing the tool that creates self-executable zip files
+# (`<preamble> <zip> <output>`). It is optional: when no toolchain is
+# registered, rules fall back to the Python `//tools/zipapp:exe_zip_maker`.
+# For now, only rules_python itself registers one (a Rust build, dev-only).
+toolchain_type(
+    name = "exe_zip_maker",
+)

python/private/zipapp/BUILD.bazel

@@ -33,6 +33,7 @@ bzl_library(
     name = "py_zipapp_rule",
     srcs = ["py_zipapp_rule.bzl"],
     deps = [
+        ":py_exe_zip_maker_toolchain",
         "//python/private:attributes",
         "//python/private:builders",
         "//python/private:common",
@@ -46,3 +47,13 @@ bzl_library(
         "@rules_python_internal//:rules_python_config",
     ],
 )
+
+bzl_library(
+    name = "py_exe_zip_maker_toolchain",
+    srcs = ["py_exe_zip_maker_toolchain.bzl"],
+    deps = [
+        "//python/private:py_interpreter_program",
+        "//python/private:sentinel_impl",
+        "//python/private:toolchain_types",
+    ],
+)

python/private/zipapp/py_exe_zip_maker_toolchain.bzl

@@ -0,0 +1,75 @@
+"""Toolchain for the tool that creates self-executable zip files."""
+
+load("//python/private:py_interpreter_program.bzl", "PyInterpreterProgramInfo")
+load("//python/private:sentinel_impl.bzl", "SentinelInfo")
+load("//python/private:toolchain_types.bzl", "EXE_ZIP_MAKER_TOOLCHAIN_TYPE")
+
+def _py_exe_zip_maker_toolchain_impl(ctx):
+    exe_zip_maker = ctx.attr.exe_zip_maker
+    if exe_zip_maker != None and SentinelInfo in exe_zip_maker:
+        exe_zip_maker = None
+
+    return [
+        platform_common.ToolchainInfo(
+            exe_zip_maker = exe_zip_maker,
+        ),
+    ]
+
+py_exe_zip_maker_toolchain = rule(
+    implementation = _py_exe_zip_maker_toolchain_impl,
+    doc = """Provides the tool used to create self-executable zip files.
+
+This provides `ToolchainInfo` with the following attributes:
+* `exe_zip_maker`: {type}`Target | None`. Invoked with three positional
+  arguments: `<preamble> <zip> <output>`. Must provide either
+  `PyInterpreterProgramInfo` or `DefaultInfo.files_to_run`. If `None`, the
+  rules fall back to their built-in implementation.
+""",
+    attrs = {
+        "exe_zip_maker": attr.label(
+            # NOTE: This is an executable, but can't use `executable = True`
+            # because the `//python:none` sentinel isn't executable.
+            # `allow_files = True` (not `allow_single_file`) because binary
+            # targets can have multiple files in `DefaultInfo.files`.
+            allow_files = True,
+            cfg = "exec",
+            doc = """
+The tool to create self-executable zip files.
+
+To indicate no tool, specify the special target {obj}`//python:none`.
+""",
+        ),
+    },
+)
+
+def get_exe_zip_maker(ctx):
+    """Returns how to run the tool for creating self-executable zips.
+
+    Prefers the tool from the `exe_zip_maker` toolchain, if one is resolved
+    and provides it. Otherwise, falls back to the rule's `_exe_zip_maker`
+    attribute.
+
+    Args:
+        ctx: The rule context. The rule must declare
+            `EXE_ZIP_MAKER_TOOLCHAIN_TYPE` as an optional toolchain and have
+            an `_exe_zip_maker` attribute.
+
+    Returns:
+        {type}`struct` with fields:
+        * `executable`: {type}`Target` to pass to `actions_run()`.
+        * `toolchain`: {type}`Label | None` the toolchain type to pass to
+          `actions_run()` so the action runs on the exec platform the tool
+          was built for. `None` when the tool is a `py_interpreter_program`
+          (`actions_run` then uses the exec tools toolchain) or when falling
+          back to the attribute.
+    """
+    toolchain = ctx.toolchains[EXE_ZIP_MAKER_TOOLCHAIN_TYPE]
+    if toolchain and toolchain.exe_zip_maker:
+        executable = toolchain.exe_zip_maker
+        if PyInterpreterProgramInfo in executable:
+            return struct(executable = executable, toolchain = None)
+        return struct(
+            executable = executable,
+            toolchain = EXE_ZIP_MAKER_TOOLCHAIN_TYPE,
+        )
+    return struct(executable = ctx.attr._exe_zip_maker, toolchain = None)

python/private/zipapp/py_zipapp_rule.bzl

@@ -18,8 +18,9 @@ load("//python/private:common_labels.bzl", "labels")
 load("//python/private:py_executable_info.bzl", "PyExecutableInfo")
 load("//python/private:py_internal.bzl", "py_internal")
 load("//python/private:py_runtime_info.bzl", "PyRuntimeInfo")
-load("//python/private:toolchain_types.bzl", "EXEC_TOOLS_TOOLCHAIN_TYPE", "LAUNCHER_MAKER_TOOLCHAIN_TYPE")
+load("//python/private:toolchain_types.bzl", "EXEC_TOOLS_TOOLCHAIN_TYPE", "EXE_ZIP_MAKER_TOOLCHAIN_TYPE", "LAUNCHER_MAKER_TOOLCHAIN_TYPE")
 load("//python/private:transition_labels.bzl", "TRANSITION_LABELS")
+load(":py_exe_zip_maker_toolchain.bzl", "get_exe_zip_maker")
 
 def _is_symlink(f):
     if hasattr(f, "is_symlink"):
@@ -225,9 +226,11 @@ def _create_self_executable_zip(ctx, preamble, zip_file):
     args.add(preamble)
     args.add(zip_file)
     args.add(pyz)
+    exe_zip_maker = get_exe_zip_maker(ctx)
     actions_run(
         ctx,
-        executable = ctx.attr._exe_zip_maker,
+        executable = exe_zip_maker.executable,
+        toolchain = exe_zip_maker.toolchain,
         arguments = [args],
         inputs = depset([preamble, zip_file]),
         outputs = [pyz],
@@ -396,7 +399,10 @@ Whether the output should be an executable zip file.
     ),
 } if not rp_config.bazel_9_or_later else {})
 
-_TOOLCHAINS = [EXEC_TOOLS_TOOLCHAIN_TYPE] + ([LAUNCHER_MAKER_TOOLCHAIN_TYPE] if rp_config.bazel_9_or_later else [])
+_TOOLCHAINS = [
+    EXEC_TOOLS_TOOLCHAIN_TYPE,
+    config_common.toolchain_type(EXE_ZIP_MAKER_TOOLCHAIN_TYPE, mandatory = False),
+] + ([LAUNCHER_MAKER_TOOLCHAIN_TYPE] if rp_config.bazel_9_or_later else [])
 
 _COMMON_RULE_DOC = """
 

tests/exe_zip_maker_toolchain/BUILD.bazel

@@ -0,0 +1,5 @@
+load(":exe_zip_maker_toolchain_tests.bzl", "exe_zip_maker_toolchain_test_suite")
+
+exe_zip_maker_toolchain_test_suite(
+    name = "exe_zip_maker_toolchain_tests",
+)

tests/exe_zip_maker_toolchain/exe_zip_maker_toolchain_tests.bzl

@@ -0,0 +1,89 @@
+"""Tests that the dev-only exe_zip_maker toolchain is used."""
+
+load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("@rules_testing//lib:truth.bzl", "matching")
+load("@rules_testing//lib:util.bzl", rt_util = "util")
+load("//python:py_binary.bzl", "py_binary")
+load("//python/zipapp:py_zipapp_binary.bzl", "py_zipapp_binary")
+load("//tests/support:support.bzl", "NOT_WINDOWS", "SUPPORTS_BZLMOD")
+
+_tests = []
+
+# When developing rules_python (bzlmod), MODULE.bazel registers a toolchain
+# that points to the Rust implementation, gated behind a flag. These tests
+# verify the rules pick it up when enabled, and use the Python fallback
+# otherwise.
+#
+# The dev toolchain is only registered under bzlmod; in WORKSPACE mode,
+# rules_rust is a stub and the Python fallback is used, so skip there.
+#
+# On Windows, py_zipapp uses the Bazel launcher instead of a self-executable
+# zip, so the exe_zip_maker toolchain is never used; skip there too.
+_COMPATIBLE_WITH = SUPPORTS_BZLMOD + NOT_WINDOWS
+_RUST_EXE_ZIP_MAKER_SUFFIX = "crates/exe_zip_maker/exe_zip_maker"
+_PYTHON_EXE_ZIP_MAKER_SUFFIX = "tools/zipapp/exe_zip_maker_.py"
+_USE_RUST_FLAG = str(Label("//dev/dev_only_toolchains:use_rust_exe_zip_maker"))
+
+def _define_py_zipapp_subject(name):
+    rt_util.helper_target(
+        py_binary,
+        name = name + "_bin",
+        srcs = ["main.py"],
+        main = "main.py",
+    )
+    rt_util.helper_target(
+        py_zipapp_binary,
+        name = name + "_subject",
+        binary = name + "_bin",
+    )
+
+def _test_py_zipapp_uses_toolchain_exe_zip_maker(name):
+    _define_py_zipapp_subject(name)
+    analysis_test(
+        name = name,
+        impl = _test_py_zipapp_uses_toolchain_exe_zip_maker_impl,
+        target = name + "_subject",
+        config_settings = {
+            _USE_RUST_FLAG: "yes",
+        },
+        attr_values = {"target_compatible_with": _COMPATIBLE_WITH},
+    )
+
+def _test_py_zipapp_uses_toolchain_exe_zip_maker_impl(env, target):
+    action = env.expect.that_target(target).action_named(
+        "PyZipAppCreateExecutableZip",
+    )
+    action.argv().contains_predicate(
+        matching.str_endswith(_RUST_EXE_ZIP_MAKER_SUFFIX),
+    )
+
+_tests.append(_test_py_zipapp_uses_toolchain_exe_zip_maker)
+
+def _test_py_zipapp_flag_disabled_uses_python_exe_zip_maker(name):
+    _define_py_zipapp_subject(name)
+    analysis_test(
+        name = name,
+        impl = _test_py_zipapp_flag_disabled_uses_python_exe_zip_maker_impl,
+        target = name + "_subject",
+        config_settings = {
+            _USE_RUST_FLAG: "no",
+        },
+        attr_values = {"target_compatible_with": _COMPATIBLE_WITH},
+    )
+
+def _test_py_zipapp_flag_disabled_uses_python_exe_zip_maker_impl(env, target):
+    action = env.expect.that_target(target).action_named(
+        "PyZipAppCreateExecutableZip",
+    )
+    action.argv().contains_predicate(
+        matching.str_endswith(_PYTHON_EXE_ZIP_MAKER_SUFFIX),
+    )
+
+_tests.append(_test_py_zipapp_flag_disabled_uses_python_exe_zip_maker)
+
+def exe_zip_maker_toolchain_test_suite(name):
+    test_suite(
+        name = name,
+        tests = _tests,
+    )

tests/exe_zip_maker_toolchain/main.py

@@ -0,0 +1 @@
+print("hello")