POC for py_test toolchain

Author: ewianda

Diff for: examples/bzlmod/BUILD.bazel

@@ -10,6 +10,7 @@ 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.
@@ -82,3 +83,23 @@ 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/test_toolchain.bzl

@@ -0,0 +1,54 @@
+"""
+Simple toolchain which overrides env and exec requirements.
+"""
+
+PytestProvider = provider(
+    fields = [
+        "get_runner",
+    ],
+)
+
+def _get_runner(ctx, binary_info):
+    executable = ctx.actions.declare_file("{}_cov".format(ctx.label.name))
+    ctx.actions.write(
+        output = executable,
+        is_executable = True,
+        content = """#!/bin/bash
+/home/ewianda/projects/rules_python/examples/bzlmod/.venv/bin/coverage run %s
+coverage_dir=$COVERAGE_DIR/pylcov.dat
+/home/ewianda/projects/rules_python/examples/bzlmod/.venv/bin/coverage lcov -o $coverage_dir
+        """ % binary_info.executable.short_path,
+    )
+    default_runfiles = ctx.runfiles(
+        files = [executable],
+    ).merge(binary_info.default_runfiles)
+    return [
+        DefaultInfo(
+            files = depset(transitive = [binary_info.files, depset([executable])]),
+            # Here is where we would return our own runner.
+            executable = executable,
+            default_runfiles = default_runfiles,
+            data_runfiles = binary_info.data_runfiles,
+        ),
+    ]
+
+def _my_cool_toolchain_impl(ctx):
+    return [platform_common.ToolchainInfo(
+        py_test_info = PytestProvider(
+            get_runner = struct(
+                func = _get_runner,
+                args = {
+                    "execution_requirements": ctx.attr.execution_requirements,
+                    "test_environment": ctx.attr.test_environment,
+                },
+            ),
+        ),
+    )]
+
+my_cool_toolchain = rule(
+    implementation = _my_cool_toolchain_impl,
+    attrs = {
+        "execution_requirements": attr.string_dict(),
+        "test_environment": attr.string_dict(),
+    },
+)

Diff for: python/BUILD.bazel

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

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

@@ -39,11 +39,12 @@ _PY_TEST_ATTRS = {
 }
 
 def _py_binary_impl(ctx):
-    return py_executable_bazel_impl(
+    _, providers = 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_executable.bzl

@@ -816,18 +816,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,
+    )
     providers = [
         DefaultInfo(
             executable = executable,
             files = default_outputs,
-            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,
-            ),
+            default_runfiles = default_runfiles,
+            data_runfiles = data_runfiles,
         ),
         create_instrumented_files_info(ctx),
         _create_run_environment_info(ctx, inherited_environment),
@@ -897,7 +905,7 @@ def _create_providers(
         runtime_details = runtime_details,
     )
     providers.extend(extra_providers)
-    return providers
+    return binary_info, providers
 
 def _create_run_environment_info(ctx, inherited_environment):
     expanded_env = {}

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

@@ -16,6 +16,7 @@
 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",
@@ -38,18 +39,28 @@ _BAZEL_PY_TEST_ATTRS = {
         executable = True,
     ),
 }
+_PY_TEST_TOOLCHAIN_TYPE = "@rules_python//python:test_runner_toolchain_type"
 
 def _py_test_impl(ctx):
-    providers = py_executable_bazel_impl(
+    binary_info, providers = py_executable_bazel_impl(
         ctx = ctx,
         is_test = True,
         inherited_environment = ctx.attr.env_inherit,
     )
     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)]),
+    },
 )