Enable custom py_binary stub_template attribute (via py_runtime)

Fixes bazelbuild#137, but unlike my approach in bazelbuild#6632, this adds an attribute to
`py_runtime` rather than to `py_binary`

Open to suggestions on the attribute name and documentation

src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java

@@ -63,8 +63,6 @@ public class BazelPythonSemantics implements PythonSemantics {
 
   public static final Runfiles.EmptyFilesSupplier GET_INIT_PY_FILES =
       new PythonUtils.GetInitPyFiles((Predicate<PathFragment> & Serializable) source -> false);
-  private static final Template STUB_TEMPLATE =
-      Template.forResource(BazelPythonSemantics.class, "python_stub_template.txt");
 
   public static final PathFragment ZIP_RUNFILES_DIRECTORY_NAME = PathFragment.create("runfiles");
 
@@ -172,12 +170,15 @@ private static void createStubFile(
       attrVersion = config.getDefaultPythonVersion();
     }
 
+    PyRuntimeInfo provider = getRuntime(ruleContext, common);
+    Artifact stubTemplateArtifact = provider.getStubScriptTemplate();
+
     // Create the stub file.
     ruleContext.registerAction(
         new TemplateExpansionAction(
             ruleContext.getActionOwner(),
+            stubTemplateArtifact,
             stubOutput,
-            STUB_TEMPLATE,
             ImmutableList.of(
                 Substitution.of("%shebang%", getStubShebang(ruleContext, common)),
                 Substitution.of(

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntime.java

@@ -110,13 +110,14 @@ public ConfiguredTarget create(RuleContext ruleContext)
       return null;
     }
     Preconditions.checkState(pythonVersion.isTargetValue());
+    Artifact stubScriptTemplate = ruleContext.getPrerequisiteArtifact("stub_script_template");
 
     PyRuntimeInfo provider =
         hermetic
             ? PyRuntimeInfo.createForInBuildRuntime(
-                interpreter, files, coverageTool, coverageFiles, pythonVersion, stubShebang)
+                interpreter, files, coverageTool, coverageFiles, pythonVersion, stubShebang, stubScriptTemplate)
             : PyRuntimeInfo.createForPlatformRuntime(
-                interpreterPath, coverageTool, coverageFiles, pythonVersion, stubShebang);
+                interpreterPath, coverageTool, coverageFiles, pythonVersion, stubShebang, stubScriptTemplate);
 
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(files)

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntimeInfo.java

@@ -32,6 +32,7 @@
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkThread;
 import net.starlark.java.syntax.Location;
+import com.google.devtools.build.lib.analysis.configuredtargets.InputFileConfiguredTarget;
 
 /**
  * Instance of the provider type that describes Python runtimes.
@@ -62,6 +63,7 @@ public final class PyRuntimeInfo implements Info, PyRuntimeInfoApi<Artifact> {
   private final PythonVersion pythonVersion;
 
   private final String stubShebang;
+  private final Artifact stubScriptTemplate;
 
   private PyRuntimeInfo(
       @Nullable Location location,
@@ -71,7 +73,8 @@ private PyRuntimeInfo(
       @Nullable Artifact coverageTool,
       @Nullable Depset coverageFiles,
       PythonVersion pythonVersion,
-      @Nullable String stubShebang) {
+      @Nullable String stubShebang,
+      Artifact stubScriptTemplate) {
     Preconditions.checkArgument((interpreterPath == null) != (interpreter == null));
     Preconditions.checkArgument((interpreter == null) == (files == null));
     Preconditions.checkArgument((coverageTool == null) == (coverageFiles == null));
@@ -88,6 +91,7 @@ private PyRuntimeInfo(
     } else {
       this.stubShebang = PyRuntimeInfoApi.DEFAULT_STUB_SHEBANG;
     }
+    this.stubScriptTemplate = stubScriptTemplate;
   }
 
   @Override
@@ -107,7 +111,8 @@ public static PyRuntimeInfo createForInBuildRuntime(
       @Nullable Artifact coverageTool,
       @Nullable NestedSet<Artifact> coverageFiles,
       PythonVersion pythonVersion,
-      @Nullable String stubShebang) {
+      @Nullable String stubShebang,
+      Artifact stubScriptTemplate) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         /*interpreterPath=*/ null,
@@ -116,7 +121,8 @@ public static PyRuntimeInfo createForInBuildRuntime(
         coverageTool,
         coverageFiles == null ? null : Depset.of(Artifact.TYPE, coverageFiles),
         pythonVersion,
-        stubShebang);
+        stubShebang,
+        stubScriptTemplate);
   }
 
   /** Constructs an instance from native rule logic (built-in location) for a platform runtime. */
@@ -125,7 +131,8 @@ public static PyRuntimeInfo createForPlatformRuntime(
       @Nullable Artifact coverageTool,
       @Nullable NestedSet<Artifact> coverageFiles,
       PythonVersion pythonVersion,
-      @Nullable String stubShebang) {
+      @Nullable String stubShebang,
+      Artifact stubScriptTemplate) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         interpreterPath,
@@ -134,7 +141,8 @@ public static PyRuntimeInfo createForPlatformRuntime(
         coverageTool,
         coverageFiles == null ? null : Depset.of(Artifact.TYPE, coverageFiles),
         pythonVersion,
-        stubShebang);
+        stubShebang,
+        stubScriptTemplate);
   }
 
   @Override
@@ -202,6 +210,11 @@ public String getStubShebang() {
     return stubShebang;
   }
 
+  @Override
+  public Artifact getStubScriptTemplate() {
+    return stubScriptTemplate;
+  }
+
   @Nullable
   public NestedSet<Artifact> getFiles() {
     try {
@@ -264,11 +277,15 @@ public PyRuntimeInfo constructor(
         Object coverageFilesUncast,
         String pythonVersion,
         String stubShebang,
+        Object stubScriptTemplateUncast,
         StarlarkThread thread)
         throws EvalException {
       String interpreterPath =
           interpreterPathUncast == NONE ? null : (String) interpreterPathUncast;
       Artifact interpreter = interpreterUncast == NONE ? null : (Artifact) interpreterUncast;
+      // TODO
+      InputFileConfiguredTarget x = (InputFileConfiguredTarget) stubScriptTemplateUncast;
+      Artifact stubScriptTemplate = x.getArtifact();
       Depset filesDepset = null;
       if (filesUncast != NONE) {
         // Validate type of filesDepset.
@@ -312,7 +329,8 @@ public PyRuntimeInfo constructor(
             coverageTool,
             coverageDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            stubScriptTemplate);
       } else {
         return new PyRuntimeInfo(
             loc,
@@ -322,7 +340,8 @@ public PyRuntimeInfo constructor(
             coverageTool,
             coverageDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            stubScriptTemplate);
       }
     }
   }

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntimeRule.java

@@ -93,6 +93,7 @@ <p>The entry point for the tool must be loadable by a python interpreter (e.g. a
         <p>Does not apply to Windows.
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("stub_shebang", STRING).value(PyRuntimeInfo.DEFAULT_STUB_SHEBANG))
+        .add(attr("stub_script_template", LABEL).allowedFileTypes(FileTypeSet.ANY_FILE).singleArtifact())
         .add(attr("output_licenses", LICENSE))
         .build();
   }

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/python/PyRuntimeInfoApi.java

@@ -17,6 +17,7 @@
 import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.docgen.annot.StarlarkConstructor;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.starlarkbuildapi.FileApi;
 import com.google.devtools.build.lib.starlarkbuildapi.core.ProviderApi;
 import javax.annotation.Nullable;
@@ -119,6 +120,15 @@ public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
               + "to Windows.")
   String getStubShebang();
 
+  @StarlarkMethod(
+      name = "stub_script_template",
+      structField = true,
+      doc =
+          "The stub script template file to use. Should have %python_binary%, "
+             + "%workspace_name%, %main%, and %imports%. See "
+             + "@bazel_tools//tools/python:python_stub_template.txt for more variables.")
+  FileT getStubScriptTemplate();
+
   /** Provider type for {@link PyRuntimeInfoApi} objects. */
   @StarlarkBuiltin(name = "Provider", documented = false, doc = "")
   interface PyRuntimeInfoProviderApi extends ProviderApi {
@@ -204,6 +214,14 @@ interface PyRuntimeInfoProviderApi extends ProviderApi {
                       + "Default is <code>"
                       + DEFAULT_STUB_SHEBANG
                       + "</code>."),
+          @Param(
+              name = "stub_script_template",
+              // allowedTypes = {
+              //   @ParamType(type = FileApi.class),
+              // },
+              positional = false,
+              named = true,
+              doc = ""),
         },
         useStarlarkThread = true,
         selfCall = true)
@@ -216,6 +234,7 @@ PyRuntimeInfoApi<?> constructor(
         Object coverageFilesUncast,
         String pythonVersion,
         String stubShebang,
+        Object stubScriptTemplate,
         StarlarkThread thread)
         throws EvalException;
   }

src/main/starlark/builtins_bzl/BUILD

@@ -21,6 +21,7 @@ genrule(
     cmd = "$(location //src:zip_builtins)" +
           " ''" +  # system zip
           " $@ src/main/starlark/builtins_bzl $(SRCS)",
+    message = "Building builtins_bzl.zip",
     output_to_bindir = 1,
     tools = ["//src:zip_builtins"],
     visibility = [

src/main/starlark/builtins_bzl/common/python/providers.bzl

@@ -14,6 +14,7 @@
 """Providers for Python rules."""
 
 DEFAULT_STUB_SHEBANG = "#!/usr/bin/env python3"
+DEFAULT_STUB_SCRIPT_TEMPLATE = "@bazel_tools//tools/python:python_stub_template.txt"
 _PYTHON_VERSION_VALUES = ["PY2", "PY3"]
 
 def _PyRuntimeInfo_init(
@@ -24,7 +25,8 @@ def _PyRuntimeInfo_init(
         coverage_tool = None,
         coverage_files = None,
         python_version,
-        stub_shebang = None):
+        stub_shebang = None,
+        stub_script_template = None):
     if (interpreter_path == None) == (interpreter == None):
         fail("exactly one of interpreter_path or interpreter must be set")
     if (interpreter == None) != (files == None):
@@ -46,6 +48,7 @@ def _PyRuntimeInfo_init(
         "coverage_files": coverage_files,
         "python_version": python_version,
         "stub_shebang": stub_shebang,
+        "stub_script_template": stub_script_template,
     }
 
 PyRuntimeInfo, _unused_raw_py_runtime_info_ctor = provider(
@@ -96,6 +99,11 @@ the same conventions as the standard CPython interpreter.
             "script used when executing `py_binary` targets.  Does not " +
             "apply to Windows."
         ),
+        "stub_script_template": (
+            "The stub script template file to use. Should have %python_binary%, " +
+            "%workspace_name%, %main%, and %imports%. See " +
+            "@bazel_tools//tools/python:python_stub_template.txt for more variables."
+        ),
     },
 )
 

src/main/starlark/builtins_bzl/common/python/py_runtime_rule.bzl

@@ -14,7 +14,7 @@
 """Implementation of py_runtime rule."""
 
 load(":common/paths.bzl", "paths")
-load(":common/python/providers.bzl", "DEFAULT_STUB_SHEBANG")
+load(":common/python/providers.bzl", "DEFAULT_STUB_SHEBANG", "DEFAULT_STUB_SCRIPT_TEMPLATE")
 
 _PyRuntimeInfo = _builtins.toplevel.PyRuntimeInfo
 
@@ -77,6 +77,7 @@ def _py_runtime_impl(ctx):
             coverage_files = coverage_files,
             python_version = python_version,
             stub_shebang = ctx.attr.stub_shebang,
+            stub_script_template = ctx.attr.stub_script_template,
         ),
         DefaultInfo(
             files = runtime_files,
@@ -179,6 +180,16 @@ See https://github.com/bazelbuild/bazel/issues/8685 for
 motivation.
 
 Does not apply to Windows.
+""",
+        ),
+        "stub_script_template": attr.label(
+            allow_single_file = True,
+            default = DEFAULT_STUB_SCRIPT_TEMPLATE,
+            doc = """
+The stub script template file to use. Should have %python_binary%,
+%workspace_name%, %main%, and %imports%.
+
+See @bazel_tools//tools/python:python_stub_template.txt for more variables.
 """,
         ),
     },

tools/python/BUILD.tools

@@ -24,6 +24,7 @@ package(default_visibility = ["//visibility:public"])
 # don't have access to Skylib here. See
 # https://github.com/bazelbuild/skydoc/issues/166.
 exports_files([
+    "python_stub_template.txt",
     "python_version.bzl",
     "srcs_version.bzl",
     "toolchain.bzl",