POC

Author: ewianda

Diff for: examples/bzlmod/.coveragerc

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

Diff for: examples/bzlmod/BUILD.bazel

@@ -10,7 +10,6 @@ load("@pip//:requirements.bzl", "all_data_requirements", "all_requirements", "al
 load("@python_3_9//:defs.bzl", py_test_with_transition = "py_test")
 load("@python_versions//3.10:defs.bzl", compile_pip_requirements_3_10 = "compile_pip_requirements")
 load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
-load(":test_toolchain.bzl", "my_cool_toolchain")
 
 # This stanza calls a rule that generates targets for managing pip dependencies
 # with pip-compile for a particular python version.
@@ -83,23 +82,3 @@ build_test(
     name = "all_requirements",
     targets = all_requirements,
 )
-# BUILD file
-
-my_cool_toolchain(
-    name = "cool_prod_linux_runner",
-)
-
-toolchain(
-    name = "prod_linux_toolchain",
-    exec_compatible_with = [
-        "@platforms//os:linux",
-        "@platforms//cpu:x86_64",
-    ],
-    target_compatible_with = [
-        "@platforms//os:linux",
-        "@platforms//cpu:x86_64",
-    ],
-    toolchain = ":cool_prod_linux_runner",
-    # TODO: Where are these for Bazel?
-    toolchain_type = "@rules_python//python:test_runner_toolchain_type",
-)

Diff for: 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",
+)
 
 # 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.

Diff for: examples/bzlmod/test_toolchain.bzl

@@ -342,6 +342,6 @@ current_py_toolchain(
 )
 
 toolchain_type(
-    name = "test_runner_toolchain_type",
+    name = "py_test_toolchain_type",
     visibility = ["//visibility:public"],
 )

Diff for: python/BUILD.bazel

@@ -39,12 +39,11 @@ _PY_TEST_ATTRS = {
 }
 
 def _py_binary_impl(ctx):
-    _, providers = py_executable_bazel_impl(
+    return py_executable_bazel_impl(
         ctx = ctx,
         is_test = False,
         inherited_environment = [],
     )
-    return providers
 
 py_binary = create_executable_rule(
     implementation = _py_binary_impl,

Diff for: python/private/common/py_binary_rule_bazel.bzl

@@ -140,7 +140,7 @@ Valid values are:
     allow_none = True,
 )
 
-def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment = []):
+def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment = [], coverage_rc):
     """Base rule implementation for a Python executable.
 
     Google and Bazel call this common base and apply customizations using the
@@ -246,6 +246,7 @@ def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment =
         inherited_environment = inherited_environment,
         semantics = semantics,
         output_groups = exec_result.output_groups,
+        coverage_rc = coverage_rc,
     )
 
 def _get_build_info(ctx, cc_toolchain):
@@ -789,6 +790,7 @@ def _create_providers(
         inherited_environment,
         runtime_details,
         output_groups,
+        coverage_rc,
         semantics):
     """Creates the providers an executable should return.
 
@@ -816,29 +818,26 @@ def _create_providers(
     Returns:
         A list of modern providers.
     """
-    default_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles(
-        ctx,
-        runfiles_details.default_runfiles,
-    )
-    data_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles(
-        ctx,
-        runfiles_details.data_runfiles,
-    )
-    binary_info = struct(
-        files = default_outputs,
-        default_runfiles = default_runfiles,
-        data_runfiles = data_runfiles,
-        executable = executable,
-    )
+    if 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()))
+    else:
+        extra_test_env = {}
     providers = [
         DefaultInfo(
             executable = executable,
             files = default_outputs,
-            default_runfiles = default_runfiles,
-            data_runfiles = data_runfiles,
+            default_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles(
+                ctx,
+                runfiles_details.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,
@@ -905,9 +904,9 @@ def _create_providers(
         runtime_details = runtime_details,
     )
     providers.extend(extra_providers)
-    return binary_info, 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(
@@ -916,6 +915,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,

Diff for: python/private/common/py_executable.bzl

@@ -101,13 +101,14 @@ def create_executable_rule(*, attrs, **kwargs):
         **kwargs
     )
 
-def py_executable_bazel_impl(ctx, *, is_test, inherited_environment):
+def py_executable_bazel_impl(ctx, *, is_test, inherited_environment, coverage_rc = None):
     """Common code for executables for Bazel."""
     return py_executable_base_impl(
         ctx = ctx,
         semantics = create_binary_semantics_bazel(),
         is_test = is_test,
         inherited_environment = inherited_environment,
+        coverage_rc = coverage_rc,
     )
 
 def create_binary_semantics_bazel():

Diff for: python/private/common/py_executable_bazel.bzl

@@ -16,12 +16,12 @@
 load("@bazel_skylib//lib:dicts.bzl", "dicts")
 load(":attributes.bzl", "AGNOSTIC_TEST_ATTRS")
 load(":common.bzl", "maybe_add_test_execution_info")
-load("//python/private:py_executable_info.bzl", "PyExecutableInfo")
 load(
     ":py_executable_bazel.bzl",
     "create_executable_rule",
     "py_executable_bazel_impl",
 )
+load("//python/private:toolchain_types.bzl", "PY_TEST_TOOLCHAIN_TYPE")
 
 _BAZEL_PY_TEST_ATTRS = {
     # This *might* be a magic attribute to help C++ coverage work. There's no
@@ -39,28 +39,28 @@ _BAZEL_PY_TEST_ATTRS = {
         executable = True,
     ),
 }
-_PY_TEST_TOOLCHAIN_TYPE = "@rules_python//python:test_runner_toolchain_type"
 
 def _py_test_impl(ctx):
-    binary_info, providers = py_executable_bazel_impl(
+    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
+    else:
+        coverage_rc = None
+    providers = py_executable_bazel_impl(
         ctx = ctx,
         is_test = True,
         inherited_environment = ctx.attr.env_inherit,
+        coverage_rc = coverage_rc,
     )
     maybe_add_test_execution_info(providers, ctx)
-    py_test_toolchain = ctx.exec_groups["test"].toolchains[_PY_TEST_TOOLCHAIN_TYPE]
-    if not py_test_toolchain:
-        return providers
-    _ = providers.pop(0)
-    test_providers = py_test_toolchain.py_test_info.get_runner.func(ctx, binary_info)
-    providers.extend(test_providers)
+
     return providers
 
 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)]),
+        "test": exec_group(toolchains = [config_common.toolchain_type(PY_TEST_TOOLCHAIN_TYPE, mandatory = False)]),
     },
 )

Diff for: python/private/common/py_test_rule_bazel.bzl

@@ -0,0 +1,27 @@
+"""
+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,
+        ),
+    },
+)

Diff for: python/private/py_test_toolchain.bzl

@@ -20,8 +20,10 @@ load(
     ":toolchain_types.bzl",
     "EXEC_TOOLS_TOOLCHAIN_TYPE",
     "PY_CC_TOOLCHAIN_TYPE",
+    "PY_TEST_TOOLCHAIN_TYPE",
     "TARGET_TOOLCHAIN_TYPE",
 )
+load(":py_test_toolchain.bzl", "py_test_toolchain")
 
 _IS_EXEC_TOOLCHAIN_ENABLED = Label("//python/config_settings:is_exec_tools_toolchain_enabled")
 
@@ -177,3 +179,15 @@ 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):
+    # Need to create a repository rule for this to work.
+    py_test_toolchain(
+        name = "_{}_py_test_toolchain".format(name),
+        coverage_rc = coverage_rc,
+    )
+    native.toolchain(
+        name = "{}_py_test_toolchain".format(name),
+        toolchain = ":_{}_py_test_toolchain".format(name),
+        toolchain_type = PY_TEST_TOOLCHAIN_TYPE,
+    )

Diff for: python/private/py_toolchain_suite.bzl

@@ -17,6 +17,7 @@
 load("@bazel_features//:features.bzl", "bazel_features")
 load("//python:repositories.bzl", "python_register_toolchains")
 load("//python:versions.bzl", "MINOR_MAPPING", "TOOL_VERSIONS")
+load("//python/private:py_toolchain_suite.bzl", "register_py_test_toolchain")
 load(":full_version.bzl", "full_version")
 load(":pythons_hub.bzl", "hub_repo")
 load(":repo_utils.bzl", "repo_utils")
@@ -74,6 +75,11 @@ def parse_modules(module_ctx):
         ignore_root_user_error = False
 
     for mod in module_ctx.modules:
+        for tag in mod.tags.converage:
+            register_py_test_toolchain(
+                name = tag.name,
+                coverage_rc = tag.coveragerc,
+            )
         module_toolchain_versions = []
 
         toolchain_attr_structs = _create_toolchain_attr_structs(mod)
@@ -401,6 +407,19 @@ can result in spurious build failures.
                 ),
             },
         ),
+        "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,
+                ),
+            },
+        ),
     },
     **_get_bazel_version_specific_kwargs()
 )