Enable custom py_binary stub_template attribute (via py_runtime)

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

Closes #16806.

PiperOrigin-RevId: 500248760
Change-Id: Ic39b0e9f8ae8063c3dedd1f67feece2e69de6306

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

@@ -36,6 +36,7 @@
 import com.google.devtools.build.lib.rules.python.PyCommon;
 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.PyRuntimeInfo;
 import com.google.devtools.build.lib.rules.python.PythonVersion;
 
 /**
@@ -221,6 +222,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

@@ -36,7 +36,6 @@
 import com.google.devtools.build.lib.analysis.actions.LauncherFileWriteAction.LaunchInfo;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.actions.Substitution;
-import com.google.devtools.build.lib.analysis.actions.Template;
 import com.google.devtools.build.lib.analysis.actions.TemplateExpansionAction;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
@@ -62,8 +61,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");
 
@@ -166,12 +163,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("%main%", common.determineMainExecutableSource()),
@@ -285,8 +284,7 @@ public void createExecutable(
 
   /** Registers an action to create a Windows Python launcher at {@code pythonLauncher}. */
   private static void createWindowsExeLauncher(
-      RuleContext ruleContext, String pythonBinary, Artifact pythonLauncher, boolean useZipFile)
-      throws InterruptedException {
+      RuleContext ruleContext, String pythonBinary, Artifact pythonLauncher, boolean useZipFile) {
     LaunchInfo launchInfo =
         LaunchInfo.builder()
             .addKeyValuePair("binary_type", "Python")
@@ -337,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();
   }
 
@@ -427,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

@@ -110,13 +110,25 @@ 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

@@ -62,6 +62,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 +72,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 +90,7 @@ private PyRuntimeInfo(
     } else {
       this.stubShebang = PyRuntimeInfoApi.DEFAULT_STUB_SHEBANG;
     }
+    this.bootstrapTemplate = bootstrapTemplate;
   }
 
   @Override
@@ -107,16 +110,18 @@ 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,
+        /* location= */ null,
+        /* interpreterPath= */ null,
         interpreter,
         Depset.of(Artifact.TYPE, files),
         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,16 +130,18 @@ 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,
+        /* location= */ null,
         interpreterPath,
-        /*interpreter=*/ null,
-        /*files=*/ null,
+        /* interpreter= */ null,
+        /* files= */ null,
         coverageTool,
         coverageFiles == null ? null : Depset.of(Artifact.TYPE, coverageFiles),
         pythonVersion,
-        stubShebang);
+        stubShebang,
+        bootstrapTemplate);
   }
 
   @Override
@@ -202,6 +209,12 @@ public String getStubShebang() {
     return stubShebang;
   }
 
+  @Override
+  @Nullable
+  public Artifact getBootstrapTemplate() {
+    return bootstrapTemplate;
+  }
+
   @Nullable
   public NestedSet<Artifact> getFiles() {
     try {
@@ -264,11 +277,16 @@ 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) {
+        bootstrapTemplate = (Artifact) bootstrapTemplateUncast;
+      }
       Depset filesDepset = null;
       if (filesUncast != NONE) {
         // Validate type of filesDepset.
@@ -306,23 +324,25 @@ public PyRuntimeInfo constructor(
         }
         return new PyRuntimeInfo(
             loc,
-            /*interpreterPath=*/ null,
+            /* interpreterPath= */ null,
             interpreter,
             filesDepset,
             coverageTool,
             coverageDepset,
             parsedPythonVersion,
-            stubShebang);
+            stubShebang,
+            bootstrapTemplate);
       } else {
         return new PyRuntimeInfo(
             loc,
             PathFragment.create(interpreterPath),
-            /*interpreter=*/ null,
-            /*files=*/ null,
+            /* interpreter= */ null,
+            /* files= */ null,
             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,16 @@ <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

@@ -46,6 +46,8 @@
 public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
 
   static final String DEFAULT_STUB_SHEBANG = "#!/usr/bin/env python3";
+  // Must call getToolsLabel() when using this.
+  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,16 @@ interface PyRuntimeInfoProviderApi extends ProviderApi {
                       + "Default is <code>"
                       + DEFAULT_STUB_SHEBANG
                       + "</code>."),
+          @Param(
+              name = "bootstrap_template",
+              allowedTypes = {
+                @ParamType(type = FileApi.class),
+                @ParamType(type = NoneType.class),
+              },
+              positional = false,
+              named = true,
+              defaultValue = "None",
+              doc = ""),
         },
         useStarlarkThread = true,
         selfCall = true)
@@ -216,6 +237,7 @@ PyRuntimeInfoApi<?> constructor(
         Object coverageFilesUncast,
         String pythonVersion,
         String stubShebang,
+        Object bootstrapTemplate,
         StarlarkThread thread)
         throws EvalException;
   }

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

@@ -13,7 +13,10 @@
 # limitations under the License.
 """Providers for Python rules."""
 
+load(":common/python/semantics.bzl", "TOOLS_REPO")
+
 DEFAULT_STUB_SHEBANG = "#!/usr/bin/env python3"
+DEFAULT_BOOTSTRAP_TEMPLATE = "@" + TOOLS_REPO + "//tools/python:python_bootstrap_template.txt"
 _PYTHON_VERSION_VALUES = ["PY2", "PY3"]
 
 def _PyRuntimeInfo_init(
@@ -24,7 +27,8 @@ def _PyRuntimeInfo_init(
         coverage_tool = None,
         coverage_files = None,
         python_version,
-        stub_shebang = None):
+        stub_shebang = None,
+        bootstrap_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 +50,7 @@ def _PyRuntimeInfo_init(
         "coverage_files": coverage_files,
         "python_version": python_version,
         "stub_shebang": stub_shebang,
+        "bootstrap_template": bootstrap_template,
     }
 
 # TODO(#15897): Rename this to PyRuntimeInfo when we're ready to replace the Java
@@ -98,6 +103,9 @@ the same conventions as the standard CPython interpreter.
             "script used when executing `py_binary` targets.  Does not " +
             "apply to Windows."
         ),
+        "bootstrap_template": (
+            "See py_runtime_rule.bzl%py_runtime.bootstrap_template for docs."
+        ),
     },
 )