Factor out Default and RuntimeInfo providers

Create partially-bound function for configuring py_test

Author: ewianda

python/private/py_binary_rule.bzl

@@ -39,11 +39,28 @@ _PY_TEST_ATTRS = {
 }
 
 def _py_binary_impl(ctx):
-    return py_executable_bazel_impl(
+    providers, binary_info, environment_info = py_executable_bazel_impl(
         ctx = ctx,
         is_test = False,
         inherited_environment = [],
     )
+    providers.extend(
+        [
+            # We construct DefaultInfo and RunEnvironmentInfo here, as other py_binary-like
+            # rules (py_test) need a different DefaultInfo and RunEnvironmentInfo.
+            DefaultInfo(
+                executable = binary_info.executable,
+                files = binary_info.files,
+                default_runfiles = binary_info.default_runfiles,
+                data_runfiles = binary_info.data_runfiles,
+            ),
+            RunEnvironmentInfo(
+                environment = environment_info.environment,
+                inherited_environment = environment_info.inherited_environment,
+            ),
+        ],
+    )
+    return providers
 
 py_binary = create_executable_rule(
     implementation = _py_binary_impl,

python/private/py_executable.bzl

@@ -59,7 +59,6 @@ load(
 load(
     ":toolchain_types.bzl",
     "EXEC_TOOLS_TOOLCHAIN_TYPE",
-    "PY_TEST_TOOLCHAIN_TYPE",
     TOOLCHAIN_TYPE = "TARGET_TOOLCHAIN_TYPE",
 )
 
@@ -179,7 +178,7 @@ def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment =
 
     imports = collect_imports(ctx, semantics)
 
-    runtime_details = _get_runtime_details(ctx, semantics, is_test)
+    runtime_details = _get_runtime_details(ctx, semantics)
     if ctx.configuration.coverage_enabled:
         extra_deps = semantics.get_coverage_deps(ctx, runtime_details)
     else:
@@ -278,7 +277,7 @@ def _declare_executable_file(ctx):
 
     return executable
 
-def _get_runtime_details(ctx, semantics, is_test):
+def _get_runtime_details(ctx, semantics):
     """Gets various information about the Python runtime to use.
 
     While most information comes from the toolchain, various legacy and
@@ -287,7 +286,6 @@ def _get_runtime_details(ctx, semantics, is_test):
     Args:
         ctx: Rule ctx
         semantics: A `BinarySemantics` struct; see `create_binary_semantics_struct`
-        is_test: bool; True if the rule is a test rule (has `test=True`), False if not
 
     Returns:
         A struct; see inline-field comments of the return value for details.
@@ -316,7 +314,6 @@ def _get_runtime_details(ctx, semantics, is_test):
         if not effective_runtime:
             fail("Unable to find Python runtime")
 
-    extra_test_env = {}
     if effective_runtime:
         direct = []  # List of files
         transitive = []  # List of depsets
@@ -329,12 +326,6 @@ def _get_runtime_details(ctx, semantics, is_test):
                 direct.append(effective_runtime.coverage_tool)
             if effective_runtime.coverage_files:
                 transitive.append(effective_runtime.coverage_files)
-        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].short_path}
-                direct.extend(coverage_rc.files.to_list())
         runtime_files = depset(direct = direct, transitive = transitive)
     else:
         runtime_files = depset()
@@ -366,9 +357,6 @@ def _get_runtime_details(ctx, semantics, is_test):
         # be included. For in-build runtimes, this shold include the interpreter
         # and any supporting files.
         runfiles = ctx.runfiles(transitive_files = runtime_files),
-        # extra_test_env: dict[str, str]; Additional environment variables to
-        # set when running the test.
-        extra_test_env = extra_test_env,
     )
 
 def _maybe_get_runtime_from_ctx(ctx):
@@ -861,22 +849,8 @@ def _create_providers(
     Returns:
         A list of modern providers.
     """
-
     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,
-            ),
-        ),
         create_instrumented_files_info(ctx),
-        _create_run_environment_info(ctx, inherited_environment, runtime_details.extra_test_env),
         PyExecutableInfo(
             main = main_py,
             runfiles_without_exe = runfiles_details.runfiles_without_exe,
@@ -946,9 +920,22 @@ def _create_providers(
         runtime_details = runtime_details,
     )
     providers.extend(extra_providers)
-    return providers
+    environemnt_info = _create_run_environment_info(ctx, inherited_environment)
+    binary_info = struct(
+        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,
+        ),
+    )
+    return providers, binary_info, environemnt_info
 
-def _create_run_environment_info(ctx, inherited_environment, extra_test_env):
+def _create_run_environment_info(ctx, inherited_environment):
     expanded_env = {}
     for key, value in ctx.attr.env.items():
         expanded_env[key] = _py_builtins.expand_location_and_make_variables(
@@ -957,8 +944,7 @@ def _create_run_environment_info(ctx, inherited_environment, extra_test_env):
             expression = value,
             targets = ctx.attr.data,
         )
-    expanded_env.update(extra_test_env)
-    return RunEnvironmentInfo(
+    return struct(
         environment = expanded_env,
         inherited_environment = inherited_environment,
     )

python/private/py_test_rule.bzl

@@ -41,13 +41,38 @@ _BAZEL_PY_TEST_ATTRS = {
 }
 
 def _py_test_impl(ctx):
-    providers = py_executable_bazel_impl(
+    providers, binary_info, environment_info = py_executable_bazel_impl(
         ctx = ctx,
         is_test = True,
         inherited_environment = ctx.attr.env_inherit,
     )
     maybe_add_test_execution_info(providers, ctx)
-    return providers
+    py_test_toolchain = ctx.exec_groups["test"].toolchains[PY_TEST_TOOLCHAIN_TYPE]
+    if py_test_toolchain:
+        py_test_info = py_test_toolchain.py_test_info
+    else:
+        providers.extend(
+            [
+                DefaultInfo(
+                    executable = binary_info.executable,
+                    files = binary_info.files,
+                    default_runfiles = binary_info.default_runfiles,
+                    data_runfiles = binary_info.data_runfiles,
+                ),
+                RunEnvironmentInfo(
+                    environment = environment_info.environment,
+                    inherited_environment = environment_info.inherited_environment,
+                ),
+            ],
+        )
+        return providers
+    test_providers = py_test_info.get_runner.func(
+        ctx,
+        binary_info,
+        environment_info,
+        **py_test_info.get_runner.args
+    )
+    return test_providers + providers
 
 py_test = create_executable_rule(
     implementation = _py_test_impl,

python/private/py_test_toolchain.bzl

@@ -25,25 +25,66 @@ load(
 PyTestProviderInfo = provider(
     doc = "Information about the pytest toolchain",
     fields = [
-        "coverage_rc",
+        "get_runner",
     ],
 )
 
+def _get_runner(ctx, binary_info, environment_info, coverage_rc):
+    """
+        Constructs and returns a list containing `DefaultInfo` and `RunEnvironmentInfo` for a test runner setup.
+
+        Args:
+            ctx: The rule context, providing access to actions, inputs, outputs, and more.
+            binary_info: A `struct` with defaultinfo details.
+                - `executable`: The executable binary.
+                - `files`: The files associated with the binary.
+                - `default_runfiles`: The default runfiles of the binary.
+                - `data_runfiles`: Additional runfiles for data dependencies.
+            environment_info: A `struct` with environment details.
+                - `environment`: A dictionary of key-value pairs for the test environment.
+                - `inherited_environment`: A list of environment variables inherited from the host.
+            coverage_rc: A `File` or `File`-like target containing coverage configuration files.
+            """
+
+    test_env = {"COVERAGE_RC": coverage_rc.files.to_list()[0].short_path}
+    test_env.update(environment_info.environment)
+
+    return [
+        DefaultInfo(
+            # Opportunity to override the executable in the binary_info with a new testrunner.
+            executable = binary_info.executable,
+            files = binary_info.files,
+            default_runfiles = binary_info.default_runfiles.merge(
+                ctx.runfiles(
+                    transitive_files = coverage_rc.files,
+                ),
+            ),
+            data_runfiles = binary_info.data_runfiles,
+        ),
+        RunEnvironmentInfo(
+            environment = test_env,
+            inherited_environment = environment_info.inherited_environment,
+        ),
+    ]
+
 def _py_test_toolchain_impl(ctx):
     return [
         platform_common.ToolchainInfo(
             py_test_info = PyTestProviderInfo(
-                coverage_rc = ctx.attr.coverage_rc,
+                get_runner = struct(
+                    func = _get_runner,
+                    args = {
+                        "coverage_rc": ctx.attr.coverage_rc,
+                    },
+                ),
             ),
         ),
     ]
 
 py_test_toolchain = rule(
     implementation = _py_test_toolchain_impl,
     attrs = {
-        "coverage_rc": attr.label(
-            allow_single_file = True,
-        ),
+        "coverage_rc": attr.label(allow_single_file = True),
     },
 )
 
@@ -94,9 +135,9 @@ py_test_toolchain_repo = repository_rule(
     doc = "Generates a toolchain hub repository",
     attrs = {
         "coverage_rc": attr.label(
-            allow_single_file = True,
             doc = "The coverage rc file",
             mandatory = True,
+            allow_single_file = True,
         ),
         "toolchain_type": attr.label(doc = "Toolchain type", mandatory = True),
     },