POC for py_test toolchain

POC

examples/bzlmod/.coveragerc

@@ -0,0 +1,6 @@
+[report]
+include_namespace_packages=True
+skip_covered = True
+[run]
+relative_files = True
+branch = True

examples/bzlmod/MODULE.bazel

@@ -36,6 +36,10 @@ python.toolchain(
     configure_coverage_tool = True,
     python_version = "3.10",
 )
+python.converage(
+    name = "coverage",
+    coveragerc = "coverage.rc",
+)
 
 # One can override the actual toolchain versions that are available, which can be useful
 # when optimizing what gets downloaded and when.

python/BUILD.bazel

@@ -363,3 +363,8 @@ exports_files([
 current_py_toolchain(
     name = "current_py_toolchain",
 )
+
+toolchain_type(
+    name = "py_test_toolchain_type",
+    visibility = ["//visibility:public"],
+)

python/private/py_executable.bzl

@@ -60,6 +60,7 @@ load(
 load(
     ":toolchain_types.bzl",
     "EXEC_TOOLS_TOOLCHAIN_TYPE",
+    "PY_TEST_TOOLCHAIN_TYPE",
     TOOLCHAIN_TYPE = "TARGET_TOOLCHAIN_TYPE",
 )
 
@@ -254,6 +255,7 @@ def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment =
         inherited_environment = inherited_environment,
         semantics = semantics,
         output_groups = exec_result.output_groups,
+        is_test = is_test,
     )
 
 def _get_build_info(ctx, cc_toolchain):
@@ -819,7 +821,8 @@ def _create_providers(
         inherited_environment,
         runtime_details,
         output_groups,
-        semantics):
+        semantics,
+        is_test):
     """Creates the providers an executable should return.
 
     Args:
@@ -851,21 +854,32 @@ def _create_providers(
     Returns:
         A list of modern providers.
     """
+
+    default_runfiles = runfiles_details.default_runfiles
+    extra_test_env = {}
+
+    if is_test:
+        py_test_toolchain = ctx.exec_groups["test"].toolchains[PY_TEST_TOOLCHAIN_TYPE]
+        if py_test_toolchain:
+            coverage_rc = py_test_toolchain.py_test_info.coverage_rc
+            extra_test_env = {"COVERAGE_RC": coverage_rc.files.to_list()[0].path}
+            default_runfiles = default_runfiles.merge(ctx.runfiles(files = coverage_rc.files.to_list()))
+
     providers = [
         DefaultInfo(
             executable = executable,
             files = default_outputs,
             default_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles(
                 ctx,
-                runfiles_details.default_runfiles,
+                default_runfiles,
             ),
             data_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles(
                 ctx,
                 runfiles_details.data_runfiles,
             ),
         ),
         create_instrumented_files_info(ctx),
-        _create_run_environment_info(ctx, inherited_environment),
+        _create_run_environment_info(ctx, inherited_environment, extra_test_env),
         PyExecutableInfo(
             main = main_py,
             runfiles_without_exe = runfiles_details.runfiles_without_exe,
@@ -937,7 +951,7 @@ def _create_providers(
     providers.extend(extra_providers)
     return providers
 
-def _create_run_environment_info(ctx, inherited_environment):
+def _create_run_environment_info(ctx, inherited_environment, extra_test_env):
     expanded_env = {}
     for key, value in ctx.attr.env.items():
         expanded_env[key] = _py_builtins.expand_location_and_make_variables(
@@ -946,6 +960,7 @@ def _create_run_environment_info(ctx, inherited_environment):
             expression = value,
             targets = ctx.attr.data,
         )
+    expanded_env.update(extra_test_env)
     return RunEnvironmentInfo(
         environment = expanded_env,
         inherited_environment = inherited_environment,

python/private/py_test_rule_bazel.bzl

@@ -14,6 +14,7 @@
 """Rule implementation of py_test for Bazel."""
 
 load("@bazel_skylib//lib:dicts.bzl", "dicts")
+load("//python/private:toolchain_types.bzl", "PY_TEST_TOOLCHAIN_TYPE")
 load(":attributes.bzl", "AGNOSTIC_TEST_ATTRS")
 load(":common.bzl", "maybe_add_test_execution_info")
 load(
@@ -52,4 +53,7 @@ py_test = create_executable_rule(
     implementation = _py_test_impl,
     attrs = dicts.add(AGNOSTIC_TEST_ATTRS, _BAZEL_PY_TEST_ATTRS),
     test = True,
+    exec_groups = {
+        "test": exec_group(toolchains = [config_common.toolchain_type(PY_TEST_TOOLCHAIN_TYPE, mandatory = False)]),
+    },
 )

python/private/py_test_toolchain.bzl

@@ -0,0 +1,77 @@
+# Copyright 2024 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.
+
+"""
+Simple toolchain which overrides env and exec requirements.
+"""
+
+PytestProvider = provider(
+    fields = [
+        "coverage_rc",
+    ],
+)
+
+def _py_test_toolchain_impl(ctx):
+    return [
+        platform_common.ToolchainInfo(
+            py_test_info = PytestProvider(
+                coverage_rc = ctx.attr.coverage_rc,
+            ),
+        ),
+    ]
+
+py_test_toolchain = rule(
+    implementation = _py_test_toolchain_impl,
+    attrs = {
+        "coverage_rc": attr.label(
+            allow_single_file = True,
+        ),
+    },
+)
+_TOOLCHAIN_TEMPLATE = """
+load(":py_test_toolchain.bzl", "py_test_toolchain")
+py_test_toolchain(
+    name = "{name}_toolchain",
+    coverage_rc = "{coverage_rc}",
+)
+
+toolchain(
+    name = "{name}",
+    target_compatible_with = [],
+    exec_compatible_with = [],
+    toolchain = "{name}_toolchain",
+    toolchain_type = "{toolchain_type}",
+)
+"""
+
+def _toolchains_repo_impl(repository_ctx):
+    build_content = _TOOLCHAIN_TEMPLATE.format(
+        name = repository_ctx.name,
+        toolchain_type = repository_ctx.toolchain_type,
+        coverage_rc = repository_ctx.attr.coverage_rc,
+    )
+    repository_ctx.file("BUILD.bazel", build_content)
+
+py_test_toolchain_repo = repository_rule(
+    _toolchains_repo_impl,
+    doc = "Generates a toolchain hub repository",
+    attrs = {
+        "toolchain_type": attr.string(doc = "Toolchain type", mandatory = True),
+        "coverage_rc": attr.label(
+            allow_single_file = True,
+            doc = "The coverage rc file",
+            mandatory = True,
+        ),
+    },
+)

python/private/py_toolchain_suite.bzl

@@ -15,11 +15,13 @@
 """Create the toolchain defs in a BUILD.bazel file."""
 
 load("@bazel_skylib//lib:selects.bzl", "selects")
+load(":py_test_toolchain.bzl", "py_test_toolchain_repo")
 load(":text_util.bzl", "render")
 load(
     ":toolchain_types.bzl",
     "EXEC_TOOLS_TOOLCHAIN_TYPE",
     "PY_CC_TOOLCHAIN_TYPE",
+    "PY_TEST_TOOLCHAIN_TYPE",
     "TARGET_TOOLCHAIN_TYPE",
 )
 
@@ -177,3 +179,13 @@ def define_local_toolchain_suites(name, version_aware_repo_names, version_unawar
             target_settings = [],
             target_compatible_with = ["@{}//:os".format(repo)],
         )
+
+def register_py_test_toolchain(name, coverage_rc, register_toolchains = True):
+    # Need to create a repository rule for this to work.
+    py_test_toolchain_repo(
+        name = "{}_py_test_toolchain".format(name),
+        coverage_rc = coverage_rc,
+        toolchain_type = str(PY_TEST_TOOLCHAIN_TYPE),
+    )
+    if register_toolchains:
+        native.toolchain(name = "{}_py_test_toolchain".format(name))

python/private/python.bzl

@@ -16,6 +16,7 @@
 
 load("@bazel_features//:features.bzl", "bazel_features")
 load("//python:versions.bzl", "DEFAULT_RELEASE_BASE_URL", "PLATFORMS", "TOOL_VERSIONS")
+load("//python/private:py_toolchain_suite.bzl", "register_py_test_toolchain")
 load(":auth.bzl", "AUTH_ATTRS")
 load(":full_version.bzl", "full_version")
 load(":python_register_toolchains.bzl", "python_register_toolchains")
@@ -80,6 +81,12 @@ def parse_modules(*, module_ctx, _fail = fail):
 
     seen_versions = {}
     for mod in module_ctx.modules:
+        for tag in mod.tags.converage:
+            register_py_test_toolchain(
+                name = tag.name,
+                coverage_rc = tag.coveragerc,
+                register_toolchains = False,
+            )
         module_toolchain_versions = []
         toolchain_attr_structs = _create_toolchain_attr_structs(
             mod = mod,
@@ -850,6 +857,19 @@ The coverage tool to be used for a particular Python interpreter. This can overr
         ),
     },
 )
+_converage = tag_class(
+    doc = """Tag class used to register Python toolchains.""",
+    attrs = {
+        "name": attr.string(
+            mandatory = True,
+            doc = "Whether or not to configure the default coverage tool for the toolchains.",
+        ),
+        "coveragerc": attr.label(
+            doc = """ """,
+            mandatory = True,
+        ),
+    },
+)
 
 python = module_extension(
     doc = """Bzlmod extension that is used to register Python toolchains.
@@ -860,6 +880,7 @@ python = module_extension(
         "single_version_override": _single_version_override,
         "single_version_platform_override": _single_version_platform_override,
         "toolchain": _toolchain,
+        "converage": _converage,
     },
     **_get_bazel_version_specific_kwargs()
 )

python/private/stage2_bootstrap_template.py

@@ -345,13 +345,19 @@ def _maybe_collect_coverage(enable):
     unique_id = uuid.uuid4()
 
     # We need for coveragepy to use relative paths.  This can only be configured
-    rcfile_name = os.path.join(coverage_dir, ".coveragerc_{}".format(unique_id))
-    with open(rcfile_name, "w") as rcfile:
-        rcfile.write(
-            """[run]
+    if os.environ.get("COVERAGE_RC"):
+        rcfile_name = os.environ["COVERAGE_RC"]
+        assert (
+            os.path.exists(rcfile_name) == True
+        ), f"Coverage rc {rcfile_name} file does not exist"
+    else:
+        rcfile_name = os.path.join(coverage_dir, ".coveragerc_{}".format(unique_id))
+        with open(rcfile_name, "w") as rcfile:
+            rcfile.write(
+                """[run]
 relative_files = True
 """
-        )
+            )
     try:
         cov = coverage.Coverage(
             config_file=rcfile_name,

python/private/toolchain_types.bzl

@@ -21,3 +21,4 @@ implementation of the toolchain.
 TARGET_TOOLCHAIN_TYPE = Label("//python:toolchain_type")
 EXEC_TOOLS_TOOLCHAIN_TYPE = Label("//python:exec_tools_toolchain_type")
 PY_CC_TOOLCHAIN_TYPE = Label("//python/cc:toolchain_type")
+PY_TEST_TOOLCHAIN_TYPE = Label("//python:py_test_toolchain_type")