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`

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

@@ -40,6 +40,7 @@
 import com.google.devtools.build.lib.rules.python.PyRuleClasses;
 import com.google.devtools.build.lib.rules.python.PyStructUtils;
 import com.google.devtools.build.lib.rules.python.PythonVersion;
+import com.google.devtools.build.lib.rules.python.PyRuntimeInfo;
 
 /**
  * Bazel-specific rule definitions for Python rules.
@@ -241,6 +242,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)))
           .addRequiredToolchains(env.getToolsLabel("//tools/python:toolchain_type"))
           .useToolchainTransition(ToolchainTransitionMode.ENABLED)
           .build();

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");
 
@@ -161,12 +160,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(
@@ -334,7 +335,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();
   }
 
@@ -424,6 +425,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

@@ -78,11 +78,12 @@ public ConfiguredTarget create(RuleContext ruleContext)
       return null;
     }
     Preconditions.checkState(pythonVersion.isTargetValue());
+    Artifact bootstrapTemplate = ruleContext.getPrerequisiteArtifact("bootstrap_template");
 
     PyRuntimeInfo provider =
         hermetic
-            ? PyRuntimeInfo.createForInBuildRuntime(interpreter, files, pythonVersion, stubShebang)
-            : PyRuntimeInfo.createForPlatformRuntime(interpreterPath, pythonVersion, stubShebang);
+            ? PyRuntimeInfo.createForInBuildRuntime(interpreter, files, pythonVersion, stubShebang, bootstrapTemplate)
+            : PyRuntimeInfo.createForPlatformRuntime(interpreterPath, 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.
@@ -60,14 +61,16 @@ 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,
       @Nullable PathFragment interpreterPath,
       @Nullable Artifact interpreter,
       @Nullable Depset files,
       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(pythonVersion.isTargetValue());
@@ -81,6 +84,7 @@ private PyRuntimeInfo(
     } else {
       this.stubShebang = PyRuntimeInfoApi.DEFAULT_STUB_SHEBANG;
     }
+    this.bootstrapTemplate = bootstrapTemplate;
   }
 
   @Override
@@ -98,26 +102,32 @@ public static PyRuntimeInfo createForInBuildRuntime(
       Artifact interpreter,
       NestedSet<Artifact> files,
       PythonVersion pythonVersion,
-      @Nullable String stubShebang) {
+      @Nullable String stubShebang,
+      @Nullable Artifact bootstrapTemplate) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         /*interpreterPath=*/ null,
         interpreter,
         Depset.of(Artifact.TYPE, files),
         pythonVersion,
-        stubShebang);
+        stubShebang,
+        bootstrapTemplate);
   }
 
   /** Constructs an instance from native rule logic (built-in location) for a platform runtime. */
   public static PyRuntimeInfo createForPlatformRuntime(
-      PathFragment interpreterPath, PythonVersion pythonVersion, @Nullable String stubShebang) {
+      PathFragment interpreterPath,
+      PythonVersion pythonVersion,
+      @Nullable String stubShebang,
+      @Nullable Artifact bootstrapTemplate) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         interpreterPath,
         /*interpreter=*/ null,
         /*files=*/ null,
         pythonVersion,
-        stubShebang);
+        stubShebang,
+        bootstrapTemplate);
   }
 
   @Override
@@ -176,6 +186,12 @@ public String getStubShebang() {
     return stubShebang;
   }
 
+  @Override
+  @Nullable
+  public Artifact getBootstrapTemplate() {
+    return bootstrapTemplate;
+  }
+
   @Nullable
   public NestedSet<Artifact> getFiles() {
     try {
@@ -215,11 +231,22 @@ public PyRuntimeInfo constructor(
         Object filesUncast,
         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.
@@ -254,15 +281,17 @@ public PyRuntimeInfo constructor(
             interpreter,
             filesDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            bootstrapTemplate);
       } else {
         return new PyRuntimeInfo(
             loc,
             PathFragment.create(interpreterPath),
             /*interpreter=*/ null,
             /*files=*/ null,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            bootstrapTemplate);
       }
     }
   }

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

@@ -68,7 +68,7 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
                 .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
@@ -77,6 +77,14 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
         <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",
@@ -99,6 +101,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 {
@@ -161,6 +172,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)
@@ -171,6 +191,7 @@ PyRuntimeInfoApi<?> constructor(
         Object filesUncast,
         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

@@ -58,7 +58,7 @@ private static void assertHasOrderAndContainsExactly(
   public void factoryMethod_InBuildRuntime() throws Exception {
     NestedSet<Artifact> files = NestedSetBuilder.create(Order.STABLE_ORDER, dummyFile);
     PyRuntimeInfo inBuildRuntime =
-        PyRuntimeInfo.createForInBuildRuntime(dummyInterpreter, files, PythonVersion.PY2, null);
+        PyRuntimeInfo.createForInBuildRuntime(dummyInterpreter, files, PythonVersion.PY2, null, dummyFile);
 
     assertThat(inBuildRuntime.getCreationLocation()).isEqualTo(Location.BUILTIN);
     assertThat(inBuildRuntime.getInterpreterPath()).isNull();
@@ -75,7 +75,7 @@ public void factoryMethod_InBuildRuntime() throws Exception {
   public void factoryMethod_PlatformRuntime() {
     PathFragment path = PathFragment.create("/system/interpreter");
     PyRuntimeInfo platformRuntime =
-        PyRuntimeInfo.createForPlatformRuntime(path, PythonVersion.PY2, null);
+        PyRuntimeInfo.createForPlatformRuntime(path, PythonVersion.PY2, null, dummyFile);
 
     assertThat(platformRuntime.getCreationLocation()).isEqualTo(Location.BUILTIN);
     assertThat(platformRuntime.getInterpreterPath()).isEqualTo(path);
@@ -95,6 +95,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");
@@ -103,6 +104,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
@@ -111,6 +113,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");
@@ -128,6 +131,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");
@@ -139,6 +143,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

@@ -48,6 +48,7 @@ write_file(
 # 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",