Windows: Implement python native launcher

Now Bazel build a Windows exe to launch the python self-extracting zip
file by default, using --windows_exe_launcher=0 to switch back to cmd
wrapper.

The extra zip file with shebang preprended is not built on Windows
anymore, even when using cmd wrapper.

Change-Id: Ic7060326f19ca6e2e73ea8d8211afd1c7618083c
PiperOrigin-RevId: 165707076

6 files changed, 79 insertions, 40 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyBinaryRule.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyBinaryRule.java
index 3fd3ca2506..74eead0f43 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyBinaryRule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyBinaryRule.java
@@ -22,6 +22,7 @@ import com.google.devtools.build.lib.analysis.BaseRuleClasses;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
 import com.google.devtools.build.lib.bazel.rules.python.BazelPyRuleClasses.PyBinaryBaseRule;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.rules.python.PythonConfiguration;
 
@@ -37,6 +38,10 @@ public final class BazelPyBinaryRule implements RuleDefinition {
     minus the extension.  For example, if your entry point is called
     <code>main.py</code>, then your name should be <code>main</code>.
     <!-- #END_BLAZE_RULE.NAME --> */
+    Label launcher = env.getLauncherLabel();
+    if (launcher != null) {
+      builder.add(attr("$launcher", LABEL).cfg(HOST).value(launcher));
+    }
     return builder
         .requiresConfigurationFragments(PythonConfiguration.class, BazelPythonConfiguration.class)
         .add(attr("$zipper", LABEL).cfg(HOST).exec().value(env.getToolsLabel("//tools/zip:zipper")))
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyTestRule.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyTestRule.java
index d69d9ae9e3..a1c409ad34 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyTestRule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPyTestRule.java
@@ -24,6 +24,7 @@ import com.google.devtools.build.lib.analysis.BaseRuleClasses;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
 import com.google.devtools.build.lib.bazel.rules.python.BazelPyRuleClasses.PyBinaryBaseRule;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleClass.Builder.RuleClassType;
 import com.google.devtools.build.lib.packages.TriState;
@@ -35,6 +36,10 @@ import com.google.devtools.build.lib.rules.python.PythonConfiguration;
 public final class BazelPyTestRule implements RuleDefinition {
   @Override
   public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
+    Label launcher = env.getLauncherLabel();
+    if (launcher != null) {
+      builder.add(attr("$launcher", LABEL).cfg(HOST).value(launcher));
+    }
     return builder
         .requiresConfigurationFragments(PythonConfiguration.class, BazelPythonConfiguration.class)
         .add(attr("$zipper", LABEL).cfg(HOST).exec().value(env.getToolsLabel("//tools/zip:zipper")))
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java
index f0846a96f4..8ef53e5587 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java
@@ -33,6 +33,7 @@ import com.google.devtools.build.lib.analysis.actions.TemplateExpansionAction;
 import com.google.devtools.build.lib.analysis.actions.TemplateExpansionAction.Substitution;
 import com.google.devtools.build.lib.analysis.actions.TemplateExpansionAction.Template;
 import com.google.devtools.build.lib.analysis.test.InstrumentedFilesCollector.InstrumentationSpec;
+import com.google.devtools.build.lib.bazel.rules.NativeLauncherUtil;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
@@ -43,6 +44,8 @@ import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -115,11 +118,6 @@ public class BazelPythonSemantics implements PythonSemantics {
     return result;
   }
 
-  /** @return An artifact next to the executable file with ".zip" suffix */
-  public Artifact getPythonZipArtifact(RuleContext ruleContext, Artifact executable) {
-    return ruleContext.getRelatedArtifact(executable.getRootRelativePath(), ".zip");
-  }
-
   /** @return An artifact next to the executable file with ".temp" suffix */
   public Artifact getPythonTemplateMainArtifact(RuleContext ruleContext, Artifact executable) {
     return ruleContext.getRelatedArtifact(executable.getRootRelativePath(), ".temp");
@@ -153,7 +151,7 @@ public class BazelPythonSemantics implements PythonSemantics {
                       config.getImportAllRepositories() ? "True" : "False")),
               true));
     } else {
-      Artifact zipFile = getPythonZipArtifact(ruleContext, executable);
+      Artifact zipFile = common.getPythonZipArtifact(executable);
       Artifact templateMain = getPythonTemplateMainArtifact(ruleContext, executable);
       // The executable zip file will unzip itself into a tmp directory and then run from there
       ruleContext.registerAction(
@@ -171,35 +169,59 @@ public class BazelPythonSemantics implements PythonSemantics {
                       config.getImportAllRepositories() ? "True" : "False")),
               true));
 
-      ruleContext.registerAction(
-          new SpawnAction.Builder()
-              .addInput(zipFile)
-              .addOutput(executable)
-              .setShellCommand(
-                  "echo '#!/usr/bin/env python' | cat - "
-                      + zipFile.getExecPathString()
-                      + " > "
-                      + executable.getExecPathString())
-              .useDefaultShellEnvironment()
-              .setMnemonic("BuildBinary")
-              .build(ruleContext));
-
-      if (OS.getCurrent() == OS.WINDOWS) {
-        Artifact executableWrapper = common.getExecutableWrapper();
+      if (OS.getCurrent() != OS.WINDOWS) {
+        ruleContext.registerAction(
+            new SpawnAction.Builder()
+                .addInput(zipFile)
+                .addOutput(executable)
+                .setShellCommand(
+                    "echo '#!/usr/bin/env python' | cat - "
+                        + zipFile.getExecPathString()
+                        + " > "
+                        + executable.getExecPathString())
+                .useDefaultShellEnvironment()
+                .setMnemonic("BuildBinary")
+                .build(ruleContext));
+      } else {
+        if (ruleContext.getConfiguration().enableWindowsExeLauncher()) {
+          return createWindowsExeLauncher(ruleContext, pythonBinary, executable);
+        }
+
         ruleContext.registerAction(
             new TemplateExpansionAction(
                 ruleContext.getActionOwner(),
-                executableWrapper,
+                executable,
                 STUB_TEMPLATE_WINDOWS,
                 ImmutableList.of(Substitution.of("%python_path%", pythonBinary)),
                 true));
-        return executableWrapper;
+        return executable;
       }
     }
 
     return executable;
   }
 
+  private static Artifact createWindowsExeLauncher(
+      RuleContext ruleContext, String pythonBinary, Artifact pythonLauncher)
+      throws InterruptedException {
+    ByteArrayOutputStream launchInfo = new ByteArrayOutputStream();
+    try {
+      NativeLauncherUtil.writeLaunchInfo(launchInfo, "binary_type", "Python");
+      NativeLauncherUtil.writeLaunchInfo(
+          launchInfo, "workspace_name", ruleContext.getWorkspaceName());
+      NativeLauncherUtil.writeLaunchInfo(launchInfo, "python_bin_path", pythonBinary);
+
+      NativeLauncherUtil.writeDataSize(launchInfo);
+    } catch (IOException e) {
+      ruleContext.ruleError(e.getMessage());
+      throw new InterruptedException();
+    }
+
+    NativeLauncherUtil.createNativeLauncherActions(ruleContext, pythonLauncher, launchInfo);
+
+    return pythonLauncher;
+  }
+
   @Override
   public void postInitBinary(RuleContext ruleContext, RunfilesSupport runfilesSupport,
       PyCommon common) throws InterruptedException {
@@ -210,7 +232,7 @@ public class BazelPythonSemantics implements PythonSemantics {
         createPythonZipAction(
             ruleContext,
             executable,
-            getPythonZipArtifact(ruleContext, executable),
+            common.getPythonZipArtifact(executable),
             getPythonTemplateMainArtifact(ruleContext, executable),
             zipper,
             runfilesSupport);
@@ -265,7 +287,7 @@ public class BazelPythonSemantics implements PythonSemantics {
     // Read each runfile from execute path, add them into zip file at the right runfiles path.
     // Filter the executable file, cause we are building it.
     for (Artifact artifact : runfilesSupport.getRunfilesArtifactsWithoutMiddlemen()) {
-      if (!artifact.equals(executable)) {
+      if (!artifact.equals(executable) && !artifact.equals(zipFile)) {
         argv.addDynamicString(
             getZipRunfilesPath(artifact.getRunfilesPath(), workspaceName)
                 + "="
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/python_stub_template_windows.txt b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/python_stub_template_windows.txt
index 9f50c6e219..ceb7bb508c 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/python/python_stub_template_windows.txt
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/python/python_stub_template_windows.txt
@@ -17,7 +17,7 @@
 
 @SETLOCAL ENABLEEXTENSIONS
 
-@rem launcher=${$0%.cmd}
-@set launcher=%~dp0%~n0
+@rem launcher=${$0%.cmd}.zip
+@set launcher=%~dp0%~n0.zip
 
 @call "%python_path%" "%launcher%" %*
diff --git a/src/main/java/com/google/devtools/build/lib/rules/python/PyBinary.java b/src/main/java/com/google/devtools/build/lib/rules/python/PyBinary.java
index 4454b7a812..de6d6be581 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/python/PyBinary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/python/PyBinary.java
@@ -128,9 +128,6 @@ public abstract class PyBinary implements RuleConfiguredTargetFactory {
     Runfiles.Builder builder = new Runfiles.Builder(
         ruleContext.getWorkspaceName(), ruleContext.getConfiguration().legacyExternalRunfiles());
     builder.addArtifact(common.getExecutable());
-    if (common.getExecutableWrapper() != null) {
-      builder.addArtifact(common.getExecutableWrapper());
-    }
     if (common.getConvertedFiles() != null) {
       builder.addSymlinks(common.getConvertedFiles());
     } else {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/python/PyCommon.java b/src/main/java/com/google/devtools/build/lib/rules/python/PyCommon.java
index f9d543f4bc..85e2f90098 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/python/PyCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/python/PyCommon.java
@@ -81,7 +81,6 @@ public final class PyCommon {
   private final RuleContext ruleContext;
 
   private Artifact executable = null;
-  private Artifact executableWrapper = null;
 
   private NestedSet<Artifact> transitivePythonSources;
 
@@ -115,10 +114,18 @@ public final class PyCommon {
     Preconditions.checkNotNull(version);
 
     validatePackageName();
-    executable = ruleContext.createOutputArtifact();
     if (OS.getCurrent() == OS.WINDOWS) {
-      executableWrapper =
-          ruleContext.getImplicitOutputArtifact(ruleContext.getTarget().getName() + ".cmd");
+      String executableSuffix;
+      if (ruleContext.getConfiguration().enableWindowsExeLauncher()) {
+        executableSuffix = ".exe";
+      } else {
+        executableSuffix = ".cmd";
+      }
+      executable =
+          ruleContext.getImplicitOutputArtifact(
+              ruleContext.getTarget().getName() + executableSuffix);
+    } else {
+      executable = ruleContext.createOutputArtifact();
     }
     if (this.version == PythonVersion.PY2AND3) {
       // TODO(bazel-team): we need to create two actions
@@ -127,9 +134,11 @@ public final class PyCommon {
 
     NestedSetBuilder<Artifact> filesToBuildBuilder =
         NestedSetBuilder.<Artifact>stableOrder().addAll(srcs).add(executable);
-    if (executableWrapper != null) {
-      filesToBuildBuilder.add(executableWrapper);
+
+    if (ruleContext.getConfiguration().buildPythonZip()) {
+      filesToBuildBuilder.add(getPythonZipArtifact(executable));
     }
+
     filesToBuild = filesToBuildBuilder.build();
 
     if (ruleContext.hasErrors()) {
@@ -139,6 +148,11 @@ public final class PyCommon {
     addPyExtraActionPseudoAction();
   }
 
+  /** @return An artifact next to the executable file with ".zip" suffix */
+  public Artifact getPythonZipArtifact(Artifact executable) {
+    return ruleContext.getRelatedArtifact(executable.getRootRelativePath(), ".zip");
+  }
+
   public void addCommonTransitiveInfoProviders(RuleConfiguredTargetBuilder builder,
       PythonSemantics semantics, NestedSet<Artifact> filesToBuild) {
 
@@ -456,10 +470,6 @@ public final class PyCommon {
     return executable;
   }
 
-  public Artifact getExecutableWrapper() {
-    return executableWrapper;
-  }
-
   public Map<PathFragment, Artifact> getConvertedFiles() {
     return convertedFiles;
   }