Create Python executable zip file

using --build_python_zip to specify it, by default it's enabled on
Windows and disabled on other platforms.

--
Change-Id: Ib992edaf70c08568816b973159a429ff7165eed8
Reviewed-on: https://bazel-review.googlesource.com/#/c/4244
MOS_MIGRATED_REVID=129326115

8 files changed, 225 insertions, 28 deletions

diff --git a/src/BUILD b/src/BUILD
index edb19cf500..e6e340e1fd 100644
--- a/src/BUILD
+++ b/src/BUILD
@@ -158,8 +158,6 @@ genrule(
             ":darwin_tools",
             "//third_party/ijar:zipper",
         ],
-        # Windows specifically exclude zipper for now.
-        ":windows": [":dummy_darwin_tools"],
         "//conditions:default": [
             ":dummy_darwin_tools",
             "//third_party/ijar:zipper",
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
index 277552a353..d5836c6d42 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
@@ -855,6 +855,14 @@ public final class BuildConfiguration {
     )
     public TriState enableRunfiles;
 
+    @Option(
+      name = "build_python_zip",
+      defaultValue = "auto",
+      category = "undocumented",
+      help = "Build python executable zip; on on Windows, off on other platforms"
+    )
+    public TriState buildPythonZip;
+
     @Override
     public FragmentOptions getHost(boolean fallback) {
       Options host = (Options) getDefault();
@@ -2324,6 +2332,17 @@ public final class BuildConfiguration {
     }
   }
 
+  public boolean buildPythonZip() {
+    switch (options.buildPythonZip) {
+      case YES:
+        return true;
+      case NO:
+        return false;
+      default:
+        return OS.getCurrent() == OS.WINDOWS;
+    }
+  }
+
   /**
    * Collects executables defined by fragments.
    */
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 021c576c70..bf197929ba 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
@@ -14,6 +14,10 @@
 
 package com.google.devtools.build.lib.bazel.rules.python;
 
+import static com.google.devtools.build.lib.packages.Attribute.ConfigurationTransition.HOST;
+import static com.google.devtools.build.lib.packages.Attribute.attr;
+import static com.google.devtools.build.lib.packages.BuildType.LABEL;
+
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
 import com.google.devtools.build.lib.bazel.rules.BazelBaseRuleClasses;
@@ -35,6 +39,7 @@ public final class BazelPyBinaryRule implements RuleDefinition {
     <!-- #END_BLAZE_RULE.NAME --> */
     return builder
         .requiresConfigurationFragments(PythonConfiguration.class, BazelPythonConfiguration.class)
+        .add(attr("$zipper", LABEL).cfg(HOST).exec().value(env.getToolsLabel("//tools/zip:zipper")))
         .build();
   }
 
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 b3b6ce2d93..d69d9ae9e3 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
@@ -14,7 +14,9 @@
 
 package com.google.devtools.build.lib.bazel.rules.python;
 
+import static com.google.devtools.build.lib.packages.Attribute.ConfigurationTransition.HOST;
 import static com.google.devtools.build.lib.packages.Attribute.attr;
+import static com.google.devtools.build.lib.packages.BuildType.LABEL;
 import static com.google.devtools.build.lib.packages.BuildType.TRISTATE;
 import static com.google.devtools.build.lib.syntax.Type.BOOLEAN;
 
@@ -35,8 +37,11 @@ public final class BazelPyTestRule implements RuleDefinition {
   public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
     return builder
         .requiresConfigurationFragments(PythonConfiguration.class, BazelPythonConfiguration.class)
-        .override(attr("testonly", BOOLEAN).value(true)
-            .nonconfigurable("policy decision: should be consistent across configurations"))
+        .add(attr("$zipper", LABEL).cfg(HOST).exec().value(env.getToolsLabel("//tools/zip:zipper")))
+        .override(
+            attr("testonly", BOOLEAN)
+                .value(true)
+                .nonconfigurable("policy decision: should be consistent across configurations"))
         /* <!-- #BLAZE_RULE(py_test).ATTRIBUTE(stamp) -->
         See the section on <a href="${link py_binary_args}">py_binary()</a> arguments, except
         that the stamp argument is set to 0 by default for tests.
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 1e48a66bc5..04d567a397 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
@@ -14,16 +14,26 @@
 
 package com.google.devtools.build.lib.bazel.rules.python;
 
+import static java.nio.charset.StandardCharsets.ISO_8859_1;
+
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ParameterFile;
+import com.google.devtools.build.lib.analysis.FilesToRunProvider;
+import com.google.devtools.build.lib.analysis.RuleConfiguredTarget.Mode;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.Runfiles.Builder;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
+import com.google.devtools.build.lib.analysis.actions.CustomCommandLine;
+import com.google.devtools.build.lib.analysis.actions.ParameterFileWriteAction;
+import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 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.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsStore;
 import com.google.devtools.build.lib.rules.python.PyCommon;
 import com.google.devtools.build.lib.rules.python.PythonSemantics;
@@ -31,7 +41,7 @@ import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.Instr
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.vfs.PathFragment;
-
+import java.io.File;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -96,9 +106,14 @@ public class BazelPythonSemantics implements PythonSemantics {
   }
 
   @Override
-  public void createExecutable(RuleContext ruleContext, PyCommon common,
-      CcLinkParamsStore ccLinkParamsStore, NestedSet<PathFragment> imports) {
-    String main = common.determineMainExecutableSource();
+  public void createExecutable(
+      RuleContext ruleContext,
+      PyCommon common,
+      CcLinkParamsStore ccLinkParamsStore,
+      NestedSet<PathFragment> imports)
+      throws InterruptedException {
+    String main = common.determineMainExecutableSource(/*withWorkspaceName=*/ true);
+    Artifact executable = common.getExecutable();
     BazelPythonConfiguration config = ruleContext.getFragment(BazelPythonConfiguration.class);
     String pythonBinary;
 
@@ -108,20 +123,157 @@ public class BazelPythonSemantics implements PythonSemantics {
       default: throw new IllegalStateException();
     }
 
-    ruleContext.registerAction(new TemplateExpansionAction(
-        ruleContext.getActionOwner(),
-        common.getExecutable(),
-        STUB_TEMPLATE,
-        ImmutableList.of(
-            Substitution.of("%main%", main),
-            Substitution.of("%python_binary%", pythonBinary),
-            Substitution.of("%imports%", Joiner.on(":").join(imports)),
-            Substitution.of("%workspace_name%", ruleContext.getWorkspaceName())),
-        true));
+    if (!ruleContext.getConfiguration().buildPythonZip()) {
+      ruleContext.registerAction(
+          new TemplateExpansionAction(
+              ruleContext.getActionOwner(),
+              executable,
+              STUB_TEMPLATE,
+              ImmutableList.of(
+                  Substitution.of("%main%", main),
+                  Substitution.of("%python_binary%", pythonBinary),
+                  Substitution.of("%imports%", Joiner.on(":").join(imports)),
+                  Substitution.of("%workspace_name%", ruleContext.getWorkspaceName())),
+              true));
+    } else {
+      Artifact zipFile = common.getPythonZipArtifact();
+      PathFragment workspaceName = getWorkspaceNameForPythonZip(ruleContext.getWorkspaceName());
+      PathFragment defaultWorkspacename = new PathFragment(Label.DEFAULT_REPOSITORY_DIRECTORY);
+      StringBuilder importPaths = new StringBuilder();
+      importPaths.append(File.pathSeparator).append("$0/").append(workspaceName);
+      for (PathFragment path : imports) {
+        if (path.startsWith(defaultWorkspacename)) {
+          path = new PathFragment(workspaceName, path.subFragment(1, path.segmentCount()));
+        }
+        importPaths.append(File.pathSeparator).append("$0/").append(path.toString());
+      }
+      String zipHeader =
+          "#!/bin/sh\n"
+              + "export PYTHONPATH=\"$PYTHONPATH"
+              + importPaths
+              + "\"\n"
+              + "exec "
+              + pythonBinary
+              + " $0 $@\n";
+      ruleContext.registerAction(
+          new SpawnAction.Builder()
+              .addInput(zipFile)
+              .addOutput(executable)
+              .setShellCommand(
+                  "echo '"
+                      + zipHeader
+                      + "' | cat - "
+                      + zipFile.getExecPathString()
+                      + " > "
+                      + executable.getExecPathString())
+              .useDefaultShellEnvironment()
+              .setMnemonic("BuildBinary")
+              .build(ruleContext));
+    }
   }
 
   @Override
   public void postInitBinary(RuleContext ruleContext, RunfilesSupport runfilesSupport,
-      PyCommon common) {
+      PyCommon common) throws InterruptedException {
+    if (ruleContext.getConfiguration().buildPythonZip()) {
+      FilesToRunProvider zipper = ruleContext.getExecutablePrerequisite("$zipper", Mode.HOST);
+      if (!ruleContext.hasErrors()) {
+        createPythonZipAction(
+            ruleContext,
+            common.getExecutable(),
+            common.getPythonZipArtifact(),
+            common.determineMainExecutableSource(false),
+            zipper,
+            runfilesSupport);
+      }
+    }
+  }
+
+  // TODO(pcloudy): This is a temporary workaround
+  private static PathFragment getWorkspaceNameForPythonZip(String workspaceName) {
+    // Currently, the default workspace name "__main__" will causing python can't find __main__.py
+    // in executable zip file. Rename it to "main"
+    if (workspaceName.equals(Label.DEFAULT_REPOSITORY_DIRECTORY)) {
+      return new PathFragment("__default__");
+    }
+    return new PathFragment(workspaceName);
+  }
+
+  private static boolean isUnderWorkspace(PathFragment path) {
+    return !path.startsWith(Label.EXTERNAL_PACKAGE_NAME);
+  }
+
+  private static String getRunfilesPath(PathFragment path, PathFragment workspaceName) {
+    if (isUnderWorkspace(path)) {
+      // If the file is under workspace, add workspace name as prefix
+      return workspaceName.getRelative(path).toString();
+    }
+    // If the file is in external package, strip "external"
+    return path.relativeTo(Label.EXTERNAL_PACKAGE_NAME).toString();
+  }
+
+  private static String getRunfilesPath(String path, PathFragment workspaceName) {
+    return getRunfilesPath(new PathFragment(path), workspaceName);
+  }
+
+  private static void createPythonZipAction(
+      RuleContext ruleContext,
+      Artifact executable,
+      Artifact zipFile,
+      String main,
+      FilesToRunProvider zipper,
+      RunfilesSupport runfilesSupport) {
+
+    NestedSetBuilder<Artifact> inputsBuilder = NestedSetBuilder.stableOrder();
+    PathFragment workspaceName = getWorkspaceNameForPythonZip(ruleContext.getWorkspaceName());
+    CustomCommandLine.Builder argv = new CustomCommandLine.Builder();
+
+    argv.add("__main__.py=" + main);
+
+    // Creating __init__.py files under each directory
+    argv.add("__init__.py=");
+    argv.add(getRunfilesPath("__init__.py", workspaceName) + "=");
+    for (String path : runfilesSupport.getRunfiles().getEmptyFilenames()) {
+      argv.add(getRunfilesPath(path, workspaceName) + "=");
+    }
+
+    // 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.getRunfiles().getArtifacts()) {
+      if (!artifact.equals(executable)) {
+        argv.add(
+            getRunfilesPath(artifact.getExecPath(), workspaceName)
+                + "="
+                + artifact.getExecPathString());
+        inputsBuilder.add(artifact);
+      }
+    }
+
+    // zipper can only consume file list options from param file not other options,
+    // so write file list in the param file first.
+    Artifact paramFile =
+        ruleContext.getDerivedArtifact(
+            ParameterFile.derivePath(zipFile.getRootRelativePath()), zipFile.getRoot());
+
+    ruleContext.registerAction(
+        new ParameterFileWriteAction(
+            ruleContext.getActionOwner(),
+            paramFile,
+            argv.build(),
+            ParameterFile.ParameterFileType.UNQUOTED,
+            ISO_8859_1));
+
+    ruleContext.registerAction(
+        new SpawnAction.Builder()
+            .addInput(paramFile)
+            .addTransitiveInputs(inputsBuilder.build())
+            .addOutput(zipFile)
+            .setExecutable(zipper)
+            .useDefaultShellEnvironment()
+            .addArgument("cC")
+            .addArgument(zipFile.getExecPathString())
+            .addArgument("@" + paramFile.getExecPathString())
+            .setMnemonic("PythonZipper")
+            .build(ruleContext));
   }
 }
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 51e6f97aae..74abb33451 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
@@ -395,10 +395,8 @@ public final class PyCommon {
     }
   }
 
-  /**
-   * @return A String that is the full path to the main python entry point.
-   */
-  public String determineMainExecutableSource() {
+  /** @return A String that is the full path to the main python entry point. */
+  public String determineMainExecutableSource(boolean withWorkspaceName) {
     String mainSourceName;
     Rule target = ruleContext.getRule();
     boolean explicitMain = target.isAttributeValueExplicitlySpecified("main");
@@ -434,15 +432,27 @@ public final class PyCommon {
       ruleContext.attributeError("srcs", buildNoMainMatchesErrorText(explicitMain, mainSourceName));
       return null;
     }
-
+    if (!withWorkspaceName) {
+      return mainArtifact.getRunfilesPath().getPathString();
+    }
     PathFragment workspaceName = new PathFragment(
         ruleContext.getRule().getPackage().getWorkspaceName());
     return workspaceName.getRelative(mainArtifact.getRunfilesPath()).getPathString();
   }
 
+  public String determineMainExecutableSource() {
+    return determineMainExecutableSource(true);
+  }
+
   public Artifact getExecutable() {
     return executable;
   }
+  /** @return An artifact next to the executable file with ".zip" suffix */
+  public Artifact getPythonZipArtifact() {
+    PathFragment original = executable.getRootRelativePath();
+    return ruleContext.getDerivedArtifact(
+        original.replaceName(original.getBaseName() + ".zip"), executable.getRoot());
+  }
 
   public Map<PathFragment, Artifact> getConvertedFiles() {
     return convertedFiles;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/python/PythonSemantics.java b/src/main/java/com/google/devtools/build/lib/rules/python/PythonSemantics.java
index 2168a879ce..d45a8093e9 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/python/PythonSemantics.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/python/PythonSemantics.java
@@ -21,7 +21,6 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsStore;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.InstrumentationSpec;
 import com.google.devtools.build.lib.vfs.PathFragment;
-
 import java.util.Collection;
 import java.util.List;
 
@@ -68,8 +67,12 @@ public interface PythonSemantics {
    *
    * <p>This should create a generating action for {@code common.getExecutable()}.
    */
-  void createExecutable(RuleContext ruleContext, PyCommon common,
-      CcLinkParamsStore ccLinkParamsStore, NestedSet<PathFragment> imports);
+  void createExecutable(
+      RuleContext ruleContext,
+      PyCommon common,
+      CcLinkParamsStore ccLinkParamsStore,
+      NestedSet<PathFragment> imports)
+      throws InterruptedException;
 
   /**
    * Called at the end of the analysis of {@code py_binary} rules.
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/mock/BazelAnalysisMock.java b/src/test/java/com/google/devtools/build/lib/analysis/mock/BazelAnalysisMock.java
index a1b7c013da..cdef6ad311 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/mock/BazelAnalysisMock.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/mock/BazelAnalysisMock.java
@@ -45,7 +45,6 @@ import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionName;
-
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -139,6 +138,12 @@ public final class BazelAnalysisMock extends AnalysisMock {
         "package(default_visibility=['//visibility:public'])",
         "exports_files(['precompile.py'])",
         "sh_binary(name='2to3', srcs=['2to3.sh'])");
+
+    config.create(
+        "/bazel_tools_workspace/tools/zip/BUILD",
+        "package(default_visibility=['//visibility:public'])",
+        "exports_files(['precompile.py'])",
+        "cc_binary(name='zipper', srcs=['zip_main.cc'])");
     ccSupport().setup(config);
   }