feat(py_runtime): Allow py_runtime to take an executable target as the interpreter (bazel-contrib#1621)

mishazharov · rickeylev · web-flow · commit 69edec83713d · 2024-01-08T09:14:35.000Z
This PR allows `py_runtime` to accept an executable (e.g. `sh_binary`). This makes it easier to customize the interpreter binary used, as it allows intercepting invocation of the interpreter. For example, it can be used to change how the interpreter searches for dynamic libraries. Related to bazel-contrib#1612 --------- Co-authored-by: Richard Levasseur <rlevasseur@google.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -37,6 +37,9 @@ A brief description of the categories of changes:
 * (pip_install) the deprecated `pip_install` macro and related items have been
   removed.
 
+* (toolchains) `py_runtime` can now take an executable target. Note: runfiles
+  from the target are not supported yet.
+
 ### Fixed
 
 * (gazelle) The gazelle plugin helper was not working with Python toolchains 3.11
diff --git a/python/private/common/py_runtime_rule.bzl b/python/private/common/py_runtime_rule.bzl
@@ -25,7 +25,7 @@ _py_builtins = py_internal
 
 def _py_runtime_impl(ctx):
     interpreter_path = ctx.attr.interpreter_path or None  # Convert empty string to None
-    interpreter = ctx.file.interpreter
+    interpreter = ctx.attr.interpreter
     if (interpreter_path and interpreter) or (not interpreter_path and not interpreter):
         fail("exactly one of the 'interpreter' or 'interpreter_path' attributes must be specified")
 
@@ -34,12 +34,30 @@ def _py_runtime_impl(ctx):
         for t in ctx.attr.files
     ])
 
+    runfiles = ctx.runfiles()
+
     hermetic = bool(interpreter)
     if not hermetic:
         if runtime_files:
             fail("if 'interpreter_path' is given then 'files' must be empty")
         if not paths.is_absolute(interpreter_path):
             fail("interpreter_path must be an absolute path")
+    else:
+        interpreter_di = interpreter[DefaultInfo]
+
+        if interpreter_di.files_to_run and interpreter_di.files_to_run.executable:
+            interpreter = interpreter_di.files_to_run.executable
+            runfiles = runfiles.merge(interpreter_di.default_runfiles)
+
+            runtime_files = depset(transitive = [
+                interpreter_di.files,
+                interpreter_di.default_runfiles.files,
+                runtime_files,
+            ])
+        elif _is_singleton_depset(interpreter_di.files):
+            interpreter = interpreter_di.files.to_list()[0]
+        else:
+            fail("interpreter must be an executable target or must produce exactly one file.")
 
     if ctx.attr.coverage_tool:
         coverage_di = ctx.attr.coverage_tool[DefaultInfo]
@@ -88,7 +106,7 @@ def _py_runtime_impl(ctx):
         BuiltinPyRuntimeInfo(**builtin_py_runtime_info_kwargs),
         DefaultInfo(
             files = runtime_files,
-            runfiles = ctx.runfiles(),
+            runfiles = runfiles,
         ),
     ]
 
@@ -186,10 +204,28 @@ runtime. For a platform runtime this attribute must not be set.
 """,
         ),
         "interpreter": attr.label(
-            allow_single_file = True,
+            # We set `allow_files = True` to allow specifying executable
+            # targets from rules that have more than one default output,
+            # e.g. sh_binary.
+            allow_files = True,
             doc = """
-For an in-build runtime, this is the target to invoke as the interpreter. For a
-platform runtime this attribute must not be set.
+For an in-build runtime, this is the target to invoke as the interpreter. It
+can be either of:
+
+* A single file, which will be the interpreter binary. It's assumed such
+  interpreters are either self-contained single-file executables or any
+  supporting files are specified in `files`.
+* An executable target. The target's executable will be the interpreter binary.
+  Any other default outputs (`target.files`) and plain files runfiles
+  (`runfiles.files`) will be automatically included as if specified in the
+  `files` attribute.
+
+  NOTE: the runfiles of the target may not yet be properly respected/propagated
+  to consumers of the toolchain/interpreter, see
+  bazelbuild/rules_python/issues/1612
+
+For a platform runtime (i.e. `interpreter_path` being set) this attribute must
+not be set.
 """,
         ),
         "interpreter_path": attr.string(doc = """
diff --git a/tests/py_runtime/py_runtime_tests.bzl b/tests/py_runtime/py_runtime_tests.bzl
@@ -30,16 +30,20 @@ _SKIP_TEST = {
 }
 
 def _simple_binary_impl(ctx):
-    output = ctx.actions.declare_file(ctx.label.name)
-    ctx.actions.write(output, "", is_executable = True)
+    executable = ctx.actions.declare_file(ctx.label.name)
+    ctx.actions.write(executable, "", is_executable = True)
     return [DefaultInfo(
-        executable = output,
+        executable = executable,
+        files = depset([executable] + ctx.files.extra_default_outputs),
         runfiles = ctx.runfiles(ctx.files.data),
     )]
 
 _simple_binary = rule(
     implementation = _simple_binary_impl,
-    attrs = {"data": attr.label_list(allow_files = True)},
+    attrs = {
+        "data": attr.label_list(allow_files = True),
+        "extra_default_outputs": attr.label_list(allow_files = True),
+    },
     executable = True,
 )
 
@@ -239,6 +243,95 @@ def _test_in_build_interpreter_impl(env, target):
 
 _tests.append(_test_in_build_interpreter)
 
+def _test_interpreter_binary_with_multiple_outputs(name):
+    rt_util.helper_target(
+        _simple_binary,
+        name = name + "_built_interpreter",
+        extra_default_outputs = ["extra_default_output.txt"],
+        data = ["runfile.txt"],
+    )
+
+    rt_util.helper_target(
+        py_runtime,
+        name = name + "_subject",
+        interpreter = name + "_built_interpreter",
+        python_version = "PY3",
+    )
+    analysis_test(
+        name = name,
+        target = name + "_subject",
+        impl = _test_interpreter_binary_with_multiple_outputs_impl,
+    )
+
+def _test_interpreter_binary_with_multiple_outputs_impl(env, target):
+    target = env.expect.that_target(target)
+    py_runtime_info = target.provider(
+        PyRuntimeInfo,
+        factory = py_runtime_info_subject,
+    )
+    py_runtime_info.interpreter().short_path_equals("{package}/{test_name}_built_interpreter")
+    py_runtime_info.files().contains_exactly([
+        "{package}/extra_default_output.txt",
+        "{package}/runfile.txt",
+        "{package}/{test_name}_built_interpreter",
+    ])
+
+    target.default_outputs().contains_exactly([
+        "{package}/extra_default_output.txt",
+        "{package}/runfile.txt",
+        "{package}/{test_name}_built_interpreter",
+    ])
+
+    target.runfiles().contains_exactly([
+        "{workspace}/{package}/runfile.txt",
+        "{workspace}/{package}/{test_name}_built_interpreter",
+    ])
+
+_tests.append(_test_interpreter_binary_with_multiple_outputs)
+
+def _test_interpreter_binary_with_single_output_and_runfiles(name):
+    rt_util.helper_target(
+        _simple_binary,
+        name = name + "_built_interpreter",
+        data = ["runfile.txt"],
+    )
+
+    rt_util.helper_target(
+        py_runtime,
+        name = name + "_subject",
+        interpreter = name + "_built_interpreter",
+        python_version = "PY3",
+    )
+    analysis_test(
+        name = name,
+        target = name + "_subject",
+        impl = _test_interpreter_binary_with_single_output_and_runfiles_impl,
+    )
+
+def _test_interpreter_binary_with_single_output_and_runfiles_impl(env, target):
+    target = env.expect.that_target(target)
+    py_runtime_info = target.provider(
+        PyRuntimeInfo,
+        factory = py_runtime_info_subject,
+    )
+    py_runtime_info.interpreter().short_path_equals("{package}/{test_name}_built_interpreter")
+    py_runtime_info.files().contains_exactly([
+        "{package}/runfile.txt",
+        "{package}/{test_name}_built_interpreter",
+    ])
+
+    target.default_outputs().contains_exactly([
+        "{package}/runfile.txt",
+        "{package}/{test_name}_built_interpreter",
+    ])
+
+    target.runfiles().contains_exactly([
+        "{workspace}/{package}/runfile.txt",
+        "{workspace}/{package}/{test_name}_built_interpreter",
+    ])
+
+_tests.append(_test_interpreter_binary_with_single_output_and_runfiles)
+
 def _test_must_have_either_inbuild_or_system_interpreter(name):
     if br_util.is_bazel_6_or_higher():
         py_runtime_kwargs = {}
diff --git a/tests/py_runtime_info_subject.bzl b/tests/py_runtime_info_subject.bzl
@@ -31,6 +31,7 @@ def py_runtime_info_subject(info, *, meta):
     # buildifier: disable=uninitialized
     public = struct(
         # go/keep-sorted start
+        actual = info,
         bootstrap_template = lambda *a, **k: _py_runtime_info_subject_bootstrap_template(self, *a, **k),
         coverage_files = lambda *a, **k: _py_runtime_info_subject_coverage_files(self, *a, **k),
         coverage_tool = lambda *a, **k: _py_runtime_info_subject_coverage_tool(self, *a, **k),
