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/BUILD

@@ -46,6 +46,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/rules/cpp",
         "//src/main/java/com/google/devtools/build/lib/rules/python",
+        "//src/main/java/com/google/devtools/build/lib/util:filetype",
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/common/options",

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

@@ -37,6 +37,8 @@
 import com.google.devtools.build.lib.rules.python.PyInfo;
 import com.google.devtools.build.lib.rules.python.PyRuleClasses;
 import com.google.devtools.build.lib.rules.python.PythonVersion;
+import com.google.devtools.build.lib.rules.python.PyRuntimeInfo;
+import com.google.devtools.build.lib.util.FileTypeSet;
 
 /**
  * Bazel-specific rule definitions for Python rules.
@@ -221,6 +223,11 @@ responsible for creating (possibly empty) __init__.py files and adding them to t
           .add(
               attr("$py_toolchain_type", NODEP_LABEL)
                   .value(env.getToolsLabel("//tools/python:toolchain_type")))
+          /* Only used when no py_runtime() is available. See #7901
+          */
+          .add(
+              attr("$default_bootstrap_template", LABEL)
+                  .value(env.getToolsLabel(PyRuntimeInfo.DEFAULT_BOOTSTRAP_TEMPLATE)))
           .addToolchainTypes(
               ToolchainTypeRequirement.builder(env.getToolsLabel("//tools/python:toolchain_type"))
                   .mandatory(true)

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

@@ -23,6 +23,7 @@
 import com.google.devtools.build.lib.actions.CommandLineItem;
 import com.google.devtools.build.lib.actions.ParamFileInfo;
 import com.google.devtools.build.lib.actions.ParameterFile;
+import com.google.devtools.build.lib.analysis.config.InvalidConfigurationException;
 import com.google.devtools.build.lib.analysis.AnalysisUtils;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
@@ -63,8 +64,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 +171,14 @@ private static void createStubFile(
       attrVersion = config.getDefaultPythonVersion();
     }
 
+    Artifact bootstrapTemplate = getBootstrapTemplate(ruleContext, common);
+
     // Create the stub file.
     ruleContext.registerAction(
         new TemplateExpansionAction(
             ruleContext.getActionOwner(),
+            bootstrapTemplate,
             stubOutput,
-            STUB_TEMPLATE,
             ImmutableList.of(
                 Substitution.of("%shebang%", getStubShebang(ruleContext, common)),
                 Substitution.of(
@@ -344,7 +345,7 @@ private static String getZipRunfilesPath(
     }
     // We put the whole runfiles tree under the ZIP_RUNFILES_DIRECTORY_NAME directory, by doing this
     // , we avoid the conflict between default workspace name "__main__" and __main__.py file.
-    // Note: This name has to be the same with the one in python_stub_template.txt.
+    // Note: This name has to be the same with the one in python_bootstrap_template.txt.
     return ZIP_RUNFILES_DIRECTORY_NAME.getRelative(zipRunfilesPath).toString();
   }
 
@@ -434,6 +435,17 @@ private static PyRuntimeInfo getRuntime(RuleContext ruleContext, PyCommon common
         : ruleContext.getPrerequisite(":py_interpreter", PyRuntimeInfo.PROVIDER);
   }
 
+  private static Artifact getBootstrapTemplate(RuleContext ruleContext, PyCommon common) {
+    PyRuntimeInfo provider = getRuntime(ruleContext, common);
+    if (provider != null) {
+      Artifact bootstrapTemplate = provider.getBootstrapTemplate();
+      if (bootstrapTemplate != null) {
+        return bootstrapTemplate;
+      }
+    }
+    return ruleContext.getPrerequisiteArtifact("$default_bootstrap_template");
+  }
+
   private static void addRuntime(
       RuleContext ruleContext, PyCommon common, Runfiles.Builder builder) {
     PyRuntimeInfo provider = getRuntime(ruleContext, common);

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 bootstrapTemplate = ruleContext.getPrerequisiteArtifact("bootstrap_template");
 
     PyRuntimeInfo provider =
         hermetic
             ? PyRuntimeInfo.createForInBuildRuntime(
-                interpreter, files, coverageTool, coverageFiles, pythonVersion, stubShebang)
+                interpreter, files, coverageTool, coverageFiles, pythonVersion, stubShebang, bootstrapTemplate)
             : PyRuntimeInfo.createForPlatformRuntime(
-                interpreterPath, coverageTool, coverageFiles, pythonVersion, stubShebang);
+                interpreterPath, coverageTool, coverageFiles, pythonVersion, stubShebang, bootstrapTemplate);
 
     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;
+  @Nullable private final Artifact bootstrapTemplate;
 
   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,
+      @Nullable Artifact bootstrapTemplate) {
     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.bootstrapTemplate = bootstrapTemplate;
   }
 
   @Override
@@ -107,7 +111,8 @@ public static PyRuntimeInfo createForInBuildRuntime(
       @Nullable Artifact coverageTool,
       @Nullable NestedSet<Artifact> coverageFiles,
       PythonVersion pythonVersion,
-      @Nullable String stubShebang) {
+      @Nullable String stubShebang,
+      @Nullable Artifact bootstrapTemplate) {
     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,
+        bootstrapTemplate);
   }
 
   /** 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,
+      @Nullable Artifact bootstrapTemplate) {
     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,
+        bootstrapTemplate);
   }
 
   @Override
@@ -202,6 +210,12 @@ public String getStubShebang() {
     return stubShebang;
   }
 
+  @Override
+  @Nullable
+  public Artifact getBootstrapTemplate() {
+    return bootstrapTemplate;
+  }
+
   @Nullable
   public NestedSet<Artifact> getFiles() {
     try {
@@ -264,11 +278,22 @@ public PyRuntimeInfo constructor(
         Object coverageFilesUncast,
         String pythonVersion,
         String stubShebang,
+        Object bootstrapTemplateUncast,
         StarlarkThread thread)
         throws EvalException {
       String interpreterPath =
           interpreterPathUncast == NONE ? null : (String) interpreterPathUncast;
       Artifact interpreter = interpreterUncast == NONE ? null : (Artifact) interpreterUncast;
+      Artifact bootstrapTemplate = null;
+      if (bootstrapTemplateUncast != NONE) {
+        if (bootstrapTemplateUncast instanceof InputFileConfiguredTarget) {
+          InputFileConfiguredTarget bootstrapTarget = (InputFileConfiguredTarget) bootstrapTemplateUncast;
+          bootstrapTemplate = bootstrapTarget.getArtifact();
+        } else if (bootstrapTemplateUncast instanceof Artifact) {
+          // SourceArtifact, possibly only in tests?
+          bootstrapTemplate = (Artifact) bootstrapTemplateUncast;
+        }
+      }
       Depset filesDepset = null;
       if (filesUncast != NONE) {
         // Validate type of filesDepset.
@@ -312,7 +337,8 @@ public PyRuntimeInfo constructor(
             coverageTool,
             coverageDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            bootstrapTemplate);
       } else {
         return new PyRuntimeInfo(
             loc,
@@ -322,7 +348,8 @@ public PyRuntimeInfo constructor(
             coverageTool,
             coverageDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            bootstrapTemplate);
       }
     }
   }

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

@@ -84,7 +84,7 @@ <p>The entry point for the tool must be loadable by a python interpreter (e.g. a
                 .allowedValues(PyRuleClasses.TARGET_PYTHON_ATTR_VALUE_SET))
 
         /* <!-- #BLAZE_RULE(py_runtime).ATTRIBUTE(stub_shebang) -->
-        "Shebang" expression prepended to the bootstrapping Python stub script
+        "Shebang" expression prepended to the bootstrapping Python script
         used when executing <code>py_binary</code> targets.
 
         <p>See <a href="https://github.com/bazelbuild/bazel/issues/8685">issue 8685</a> for
@@ -93,6 +93,14 @@ <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))
+
+        /* <!-- #BLAZE_RULE(py_runtime).ATTRIBUTE(bootstrap_template) -->
+		Previously referred to as the "Python stub script", this is the
+        entrypoint to every Python executable target.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("bootstrap_template", LABEL)
+            .value(env.getToolsLabel(PyRuntimeInfo.DEFAULT_BOOTSTRAP_TEMPLATE))
+            .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;
@@ -46,6 +47,7 @@
 public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
 
   static final String DEFAULT_STUB_SHEBANG = "#!/usr/bin/env python3";
+  static final String DEFAULT_BOOTSTRAP_TEMPLATE = "//tools/python:python_bootstrap_template.txt";
 
   @StarlarkMethod(
       name = "interpreter_path",
@@ -119,6 +121,15 @@ public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
               + "to Windows.")
   String getStubShebang();
 
+  @StarlarkMethod(
+      name = "bootstrap_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_bootstrap_template.txt for more variables.")
+  FileT getBootstrapTemplate();
+
   /** Provider type for {@link PyRuntimeInfoApi} objects. */
   @StarlarkBuiltin(name = "Provider", documented = false, doc = "")
   interface PyRuntimeInfoProviderApi extends ProviderApi {
@@ -204,6 +215,15 @@ interface PyRuntimeInfoProviderApi extends ProviderApi {
                       + "Default is <code>"
                       + DEFAULT_STUB_SHEBANG
                       + "</code>."),
+          @Param(
+              name = "bootstrap_template",
+              // allowedTypes = {
+              //   @ParamType(type = FileApi.class),
+              // },
+              positional = false,
+              named = true,
+              defaultValue = "None",
+              doc = ""),
         },
         useStarlarkThread = true,
         selfCall = true)
@@ -216,6 +236,7 @@ PyRuntimeInfoApi<?> constructor(
         Object coverageFilesUncast,
         String pythonVersion,
         String stubShebang,
+        Object bootstrapTemplate,
         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/test/java/com/google/devtools/build/lib/packages/util/BazelMockPythonSupport.java

@@ -39,6 +39,7 @@ public void setup(MockToolsConfig config) throws IOException {
     addTool(config, "tools/python/toolchain.bzl");
     addTool(config, "tools/python/utils.bzl");
     addTool(config, "tools/python/private/defs.bzl");
+    addTool(config, "tools/python/python_bootstrap_template.txt");
 
     config.create(
         TestConstants.TOOLS_REPOSITORY_SCRATCH + "tools/python/BUILD",

src/test/java/com/google/devtools/build/lib/rules/python/PyRuntimeInfoTest.java

@@ -59,7 +59,7 @@ public void factoryMethod_InBuildRuntime() throws Exception {
     NestedSet<Artifact> files = NestedSetBuilder.create(Order.STABLE_ORDER, dummyFile);
     PyRuntimeInfo inBuildRuntime =
         PyRuntimeInfo.createForInBuildRuntime(
-            dummyInterpreter, files, null, null, PythonVersion.PY2, null);
+            dummyInterpreter, files, null, null, PythonVersion.PY2, null, dummyFile);
 
     assertThat(inBuildRuntime.getCreationLocation()).isEqualTo(Location.BUILTIN);
     assertThat(inBuildRuntime.getInterpreterPath()).isNull();
@@ -76,7 +76,7 @@ public void factoryMethod_InBuildRuntime() throws Exception {
   public void factoryMethod_PlatformRuntime() {
     PathFragment path = PathFragment.create("/system/interpreter");
     PyRuntimeInfo platformRuntime =
-        PyRuntimeInfo.createForPlatformRuntime(path, null, null, PythonVersion.PY2, null);
+        PyRuntimeInfo.createForPlatformRuntime(path, null, null, PythonVersion.PY2, null, dummyFile);
 
     assertThat(platformRuntime.getCreationLocation()).isEqualTo(Location.BUILTIN);
     assertThat(platformRuntime.getInterpreterPath()).isEqualTo(path);
@@ -96,6 +96,7 @@ public void starlarkConstructor_InBuildRuntime() throws Exception {
         "    interpreter = dummy_interpreter,",
         "    files = depset([dummy_file]),",
         "    python_version = 'PY2',",
+        "    bootstrap_template = dummy_file,",
         ")");
     PyRuntimeInfo info = (PyRuntimeInfo) ev.lookup("info");
     assertThat(info.getCreationLocation().toString()).isEqualTo(":1:21");
@@ -104,6 +105,7 @@ public void starlarkConstructor_InBuildRuntime() throws Exception {
     assertHasOrderAndContainsExactly(info.getFiles(), Order.STABLE_ORDER, dummyFile);
     assertThat(info.getPythonVersion()).isEqualTo(PythonVersion.PY2);
     assertThat(info.getStubShebang()).isEqualTo(PyRuntimeInfo.DEFAULT_STUB_SHEBANG);
+    assertThat(info.getBootstrapTemplate()).isEqualTo(dummyFile);
   }
 
   @Test
@@ -112,6 +114,7 @@ public void starlarkConstructor_PlatformRuntime() throws Exception {
         "info = PyRuntimeInfo(", //
         "    interpreter_path = '/system/interpreter',",
         "    python_version = 'PY2',",
+        "    bootstrap_template = dummy_file,",
         ")");
     PyRuntimeInfo info = (PyRuntimeInfo) ev.lookup("info");
     assertThat(info.getCreationLocation().toString()).isEqualTo(":1:21");
@@ -129,6 +132,7 @@ public void starlarkConstructor_CustomShebang() throws Exception {
         "    interpreter_path = '/system/interpreter',",
         "    python_version = 'PY2',",
         "    stub_shebang = '#!/usr/bin/custom',",
+        "    bootstrap_template = dummy_file,",
         ")");
     PyRuntimeInfo info = (PyRuntimeInfo) ev.lookup("info");
     assertThat(info.getStubShebang()).isEqualTo("#!/usr/bin/custom");
@@ -140,6 +144,7 @@ public void starlarkConstructor_FilesDefaultsToEmpty() throws Exception {
         "info = PyRuntimeInfo(", //
         "    interpreter = dummy_interpreter,",
         "    python_version = 'PY2',",
+        "    bootstrap_template = dummy_file,",
         ")");
     PyRuntimeInfo info = (PyRuntimeInfo) ev.lookup("info");
     assertHasOrderAndContainsExactly(info.getFiles(), Order.STABLE_ORDER);

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_bootstrap_template.txt",
     "python_version.bzl",
     "srcs_version.bzl",
     "toolchain.bzl",