Add stub_shebang to py_runtime

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

@@ -164,6 +164,9 @@ private static void createStubFile(
             stubOutput,
             STUB_TEMPLATE,
             ImmutableList.of(
+                Substitution.of(
+                    "%shebang%", getStubShebang(ruleContext, common)
+                ),
                 Substitution.of(
                     "%main%", common.determineMainExecutableSource(/*withWorkspaceName=*/ true)),
                 Substitution.of("%python_binary%", pythonBinary),
@@ -445,6 +448,15 @@ private static String getPythonBinary(
     return pythonBinary;
   }
 
+  private static String getStubShebang(RuleContext ruleContext, PyCommon common) {
+    PyRuntimeInfo provider = getRuntime(ruleContext, common);
+    if (provider != null) {
+      return provider.getStubShebang();
+    } else {
+      return PyRuntimeInfo.DEFAULT_STUB_SHEBANG;
+    }
+  }
+
   @Override
   public CcInfo buildCcInfoProvider(Iterable<? extends TransitiveInfoCollection> deps) {
     ImmutableList<CcInfo> ccInfos =

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

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+%shebang%
 
 # This script must retain compatibility with a wide variety of Python versions
 # since it is run for every py_binary target. Currently we guarantee support

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

@@ -44,6 +44,8 @@ public ConfiguredTarget create(RuleContext ruleContext) throws ActionConflictExc
     PythonVersion pythonVersion =
         PythonVersion.parseTargetOrSentinelValue(
             ruleContext.attributes().get("python_version", Type.STRING));
+    String stubShebang =
+          ruleContext.attributes().get("stub_shebang", Type.STRING);
 
     // Determine whether we're pointing to an in-build target (hermetic) or absolute system path
     // (non-hermetic).
@@ -82,8 +84,8 @@ public ConfiguredTarget create(RuleContext ruleContext) throws ActionConflictExc
 
     PyRuntimeInfo provider =
         hermetic
-            ? PyRuntimeInfo.createForInBuildRuntime(interpreter, files, pythonVersion)
-            : PyRuntimeInfo.createForPlatformRuntime(interpreterPath, pythonVersion);
+            ? PyRuntimeInfo.createForInBuildRuntime(interpreter, files, pythonVersion, stubShebang)
+            : PyRuntimeInfo.createForPlatformRuntime(interpreterPath, pythonVersion, stubShebang);
 
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(files)

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

@@ -54,17 +54,19 @@ public final class PyRuntimeInfo implements Info, PyRuntimeInfoApi<Artifact> {
   private final Location location;
   @Nullable private final PathFragment interpreterPath;
   @Nullable private final Artifact interpreter;
-  // Validated on initalization to contain Artifact
+  // Validated on initialization to contain Artifact
   @Nullable private final Depset files;
   /** Invariant: either PY2 or PY3. */
   private final PythonVersion pythonVersion;
+  private final String stubShebang;
 
   private PyRuntimeInfo(
       @Nullable Location location,
       @Nullable PathFragment interpreterPath,
       @Nullable Artifact interpreter,
       @Nullable Depset files,
-      PythonVersion pythonVersion) {
+      PythonVersion pythonVersion,
+      @Nullable String stubShebang) {
     Preconditions.checkArgument((interpreterPath == null) != (interpreter == null));
     Preconditions.checkArgument((interpreter == null) == (files == null));
     Preconditions.checkArgument(pythonVersion.isTargetValue());
@@ -73,6 +75,11 @@ private PyRuntimeInfo(
     this.interpreterPath = interpreterPath;
     this.interpreter = interpreter;
     this.pythonVersion = pythonVersion;
+    if (stubShebang != null && !stubShebang.isEmpty()) {
+      this.stubShebang = stubShebang;
+    } else {
+      this.stubShebang = PyRuntimeInfoApi.DEFAULT_STUB_SHEBANG;
+    }
   }
 
   @Override
@@ -87,20 +94,31 @@ public Location getCreationLoc() {
 
   /** Constructs an instance from native rule logic (built-in location) for an in-build runtime. */
   public static PyRuntimeInfo createForInBuildRuntime(
-      Artifact interpreter, NestedSet<Artifact> files, PythonVersion pythonVersion) {
+      Artifact interpreter,
+      NestedSet<Artifact> files,
+      PythonVersion pythonVersion,
+      @Nullable String stubShebang) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         /*interpreterPath=*/ null,
         interpreter,
         Depset.of(Artifact.TYPE, files),
-        pythonVersion);
+        pythonVersion,
+        stubShebang);
   }
 
   /** Constructs an instance from native rule logic (built-in location) for a platform runtime. */
   public static PyRuntimeInfo createForPlatformRuntime(
-      PathFragment interpreterPath, PythonVersion pythonVersion) {
+      PathFragment interpreterPath,
+      PythonVersion pythonVersion,
+      @Nullable String stubShebang) {
     return new PyRuntimeInfo(
-        /*location=*/ null, interpreterPath, /*interpreter=*/ null, /*files=*/ null, pythonVersion);
+        /*location=*/ null,
+        interpreterPath,
+        /*interpreter=*/ null,
+        /*files=*/ null,
+        pythonVersion,
+        stubShebang);
   }
 
   @Override
@@ -113,12 +131,14 @@ public boolean equals(Object other) {
     PyRuntimeInfo otherInfo = (PyRuntimeInfo) other;
     return (this.interpreterPath.equals(otherInfo.interpreterPath)
         && this.interpreter.equals(otherInfo.interpreter)
-        && this.files.equals(otherInfo.files));
+        && this.files.equals(otherInfo.files)
+        && this.stubShebang.equals(otherInfo.stubShebang));
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(PyRuntimeInfo.class, interpreterPath, interpreter, files);
+    return Objects.hash(PyRuntimeInfo.class, interpreterPath, interpreter,
+                        files, stubShebang);
   }
 
   /**
@@ -153,6 +173,11 @@ public Artifact getInterpreter() {
     return interpreter;
   }
 
+  @Override
+  public String getStubShebang() {
+    return stubShebang;
+  }
+
   @Nullable
   public NestedSet<Artifact> getFiles() {
     try {
@@ -191,6 +216,7 @@ public PyRuntimeInfo constructor(
         Object interpreterUncast,
         Object filesUncast,
         String pythonVersion,
+        String stubShebang,
         StarlarkThread thread)
         throws EvalException {
       String interpreterPath =
@@ -225,14 +251,20 @@ public PyRuntimeInfo constructor(
           filesDepset = Depset.of(Artifact.TYPE, NestedSetBuilder.emptySet(Order.STABLE_ORDER));
         }
         return new PyRuntimeInfo(
-            loc, /*interpreterPath=*/ null, interpreter, filesDepset, parsedPythonVersion);
+            loc,
+            /*interpreterPath=*/ null,
+            interpreter,
+            filesDepset,
+            parsedPythonVersion,
+            stubShebang);
       } else {
         return new PyRuntimeInfo(
             loc,
             PathFragment.create(interpreterPath),
             /*interpreter=*/ null,
             /*files=*/ null,
-            parsedPythonVersion);
+            parsedPythonVersion,
+            stubShebang);
       }
     }
   }

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

@@ -66,6 +66,17 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
             attr("python_version", STRING)
                 .value(PythonVersion._INTERNAL_SENTINEL.toString())
                 .allowedValues(PyRuleClasses.TARGET_PYTHON_ATTR_VALUE_SET))
+
+        /* <!-- #BLAZE_RULE(py_runtime).ATTRIBUTE(stub_shebang) -->
+        "Shebang" expression prepended to the bootstrapping Python stub script
+        used when executing <code>py_binary</code> targets.
+
+        <p>See <a href="https://github.com/bazelbuild/bazel/issues/8685">issue 8685</a> for
+        motivation.
+
+        <p>Does not apply to Windows.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("stub_shebang", STRING).value(PyRuntimeInfo.DEFAULT_STUB_SHEBANG))
         .add(attr("output_licenses", LICENSE))
         .build();
   }

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

@@ -43,6 +43,8 @@
     category = StarlarkDocumentationCategory.PROVIDER)
 public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
 
+  String DEFAULT_STUB_SHEBANG = "#!/usr/bin/env python";
+
   @StarlarkMethod(
       name = "interpreter_path",
       structField = true,
@@ -86,6 +88,15 @@ public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
               + "(only) <code>\"PY2\"</code> and <code>\"PY3\"</code>.")
   String getPythonVersionForStarlark();
 
+  @StarlarkMethod(
+      name = "stub_shebang",
+      structField = true,
+      doc =
+          "\"Shebang\" expression prepended to the bootstrapping Python stub script "
+              + "used when executing <code>py_binary</code> targets.  Does not apply "
+              + "to Windows.")
+  String getStubShebang();
+
   /** Provider type for {@link PyRuntimeInfoApi} objects. */
   @StarlarkBuiltin(name = "Provider", documented = false, doc = "")
   interface PyRuntimeInfoProviderApi extends ProviderApi {
@@ -133,6 +144,15 @@ interface PyRuntimeInfoProviderApi extends ProviderApi {
               positional = false,
               named = true,
               doc = "The value for the new object's <code>python_version</code> field."),
+          @Param(
+              name = "stub_shebang",
+              type = String.class,
+              positional = false,
+              named = true,
+              defaultValue = "'" + DEFAULT_STUB_SHEBANG + "'",
+              doc =
+                  "The value for the new object's <code>stub_shebang</code> field. "
+                      + "Default is <code>" + DEFAULT_STUB_SHEBANG + "</code>."),
         },
         useStarlarkThread = true,
         selfCall = true)
@@ -142,6 +162,7 @@ PyRuntimeInfoApi<?> constructor(
         Object interpreterUncast,
         Object filesUncast,
         String pythonVersion,
+        String stubShebang,
         StarlarkThread thread)
         throws EvalException;
   }

src/main/java/com/google/devtools/build/skydoc/fakebuildapi/python/FakePyRuntimeInfo.java

@@ -44,6 +44,9 @@ public String getPythonVersionForStarlark() {
     return "";
   }
 
+  @Override
+  public String getStubShebang() { return ""; }
+
   @Override
   public void repr(Printer printer) {}
 
@@ -56,6 +59,7 @@ public PyRuntimeInfoApi<?> constructor(
         Object interpreterUncast,
         Object filesUncast,
         String pythonVersion,
+        String stubShebang,
         StarlarkThread thread)
         throws EvalException {
       return new FakePyRuntimeInfo();

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

@@ -45,14 +45,14 @@ private static String join(String... lines) {
   }
 
   /**
-   * Given a {@code py_binary} or {@code py_test} target, returns the path of the Python interpreter
-   * used by the generated stub script.
+   * Given a {@code py_binary} or {@code py_test} target and substitution key, returns the
+   * corresponding substitution value used by the generated stub script.
    *
    * <p>This works by casting the stub script's generating action to a template expansion action and
-   * looking for the expansion key for the Python interpreter. It's therefore linked to the
-   * implementation of the rule, but that's the cost we pay for avoiding an execution-time test.
+   * looking for the requested substitution key. It's therefore linked to the implementation of the
+   * rule, but that's the cost we pay for avoiding an execution-time test.
    */
-  private String getInterpreterPathFromStub(ConfiguredTarget pyExecutableTarget) {
+  private String getSubstitutionValueFromStub(ConfiguredTarget pyExecutableTarget, String substitutionKey) {
     // First find the stub script. Normally this is just the executable associated with the target.
     // But for Windows the executable is a separate launcher with an ".exe" extension, and the stub
     // script artifact has the same base name with the extension ".temp" instead. (At least, when
@@ -75,12 +75,21 @@ private String getInterpreterPathFromStub(ConfiguredTarget pyExecutableTarget) {
     assertThat(generatingAction).isInstanceOf(TemplateExpansionAction.class);
     TemplateExpansionAction templateAction = (TemplateExpansionAction) generatingAction;
     for (Substitution sub : templateAction.getSubstitutions()) {
-      if (sub.getKey().equals("%python_binary%")) {
+      if (sub.getKey().equals(substitutionKey)) {
         return sub.getValue();
       }
     }
     throw new AssertionError(
-        "Failed to find the '%python_binary%' key in the stub script's template expansion action");
+        "Failed to find the '" + substitutionKey + "' key in the stub script's template "
+            + "expansion action");
+  }
+
+  private String getInterpreterPathFromStub(ConfiguredTarget pyExecutableTarget) {
+    return getSubstitutionValueFromStub(pyExecutableTarget, "%python_binary%");
+  }
+
+  private String getShebangFromStub(ConfiguredTarget pyExecutableTarget) {
+    return getSubstitutionValueFromStub(pyExecutableTarget, "%shebang%");
   }
 
   // TODO(#8169): Delete tests of the legacy --python_top / --python_path behavior.
@@ -168,11 +177,13 @@ private void defineToolchains() throws Exception {
         "    name = 'py2_runtime',",
         "    interpreter_path = '/system/python2',",
         "    python_version = 'PY2',",
+        "    stub_shebang = '#!/usr/bin/env python',",
         ")",
         "py_runtime(",
         "    name = 'py3_runtime',",
         "    interpreter_path = '/system/python3',",
         "    python_version = 'PY3',",
+        "    stub_shebang = '#!/usr/bin/env python3',",
         ")",
         "py_runtime_pair(",
         "    name = 'py_runtime_pair',",
@@ -214,10 +225,18 @@ public void runtimeObtainedFromToolchain() throws Exception {
         "--incompatible_use_python_toolchains=true",
         "--extra_toolchains=//toolchains:py_toolchain");
 
-    String py2Path = getInterpreterPathFromStub(getConfiguredTarget("//pkg:py2_bin"));
-    String py3Path = getInterpreterPathFromStub(getConfiguredTarget("//pkg:py3_bin"));
+    ConfiguredTarget py2 = getConfiguredTarget("//pkg:py2_bin");
+    ConfiguredTarget py3 = getConfiguredTarget("//pkg:py3_bin");
+
+    String py2Path = getInterpreterPathFromStub(py2);
+    String py3Path = getInterpreterPathFromStub(py3);
     assertThat(py2Path).isEqualTo("/system/python2");
     assertThat(py3Path).isEqualTo("/system/python3");
+
+    String py2Shebang = getShebangFromStub(py2);
+    String py3Shebang = getShebangFromStub(py3);
+    assertThat(py2Shebang).isEqualTo("#!/usr/bin/env python");
+    assertThat(py3Shebang).isEqualTo("#!/usr/bin/env python3");
   }
 
   @Test