Enable simple compilation of Python artifacts for py_binary and py_library.

--
MOS_MIGRATED_REVID=99163140

6 files changed, 90 insertions, 9 deletions

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 e560f6dd20..416e77c450 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
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.bazel.rules.python;
 
 import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.actions.Artifact;
 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;
@@ -27,6 +28,8 @@ import com.google.devtools.build.lib.rules.python.PythonSemantics;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.InstrumentationSpec;
 import com.google.devtools.build.lib.util.FileTypeSet;
 
+import java.util.Collection;
+
 /**
  * Functionality specific to the Python rules in Bazel.
  */
@@ -55,6 +58,12 @@ public class BazelPythonSemantics implements PythonSemantics {
   }
 
   @Override
+  public Collection<Artifact> precompiledPythonFiles(
+      RuleContext ruleContext, Collection<Artifact> sources, PyCommon common) {
+    return ImmutableList.of();
+  }
+
+  @Override
   public void createExecutable(RuleContext ruleContext, PyCommon common,
       CcLinkParamsStore ccLinkParamsStore) {
     String main = common.determineMainExecutableSource();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/python/PrecompilePythonMode.java b/src/main/java/com/google/devtools/build/lib/rules/python/PrecompilePythonMode.java
index b77473e7e8..d88e777b07 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/python/PrecompilePythonMode.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/python/PrecompilePythonMode.java
@@ -17,9 +17,10 @@ package com.google.devtools.build.lib.rules.python;
 /**
  * Enumerates the different modes of Python precompilation.
  *
- * NONE denotes that no precompilation should take place, and PROTO causes
- * only the generated _pb/_pb2.py files to be precompiled.
+ * <p>NONE denotes that no precompilation should take place, and PROTO causes
+ * only the generated _pb/_pb2.py files to be precompiled. ALL compiles all
+ * Python files.
  */
 public enum PrecompilePythonMode {
-  NONE, PROTO; // TODO(scottgw) add 'ALL' element.
+  NONE, PROTO, ALL;
 }
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 337d63b1a5..b732cae2ed 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
@@ -27,6 +27,7 @@ import com.google.devtools.build.lib.rules.cpp.CcLinkParams;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsProvider;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsStore;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -53,10 +54,13 @@ public abstract class PyBinary implements RuleConfiguredTargetFactory {
 
   static RuleConfiguredTargetBuilder init(
       RuleContext ruleContext, PythonSemantics semantics, PyCommon common) {
-    List<Artifact> srcs = common.validateSrcs();
     CcLinkParamsStore ccLinkParamsStore = initializeCcLinkParamStore(ruleContext);
 
-    common.initBinary(srcs);
+    List<Artifact> srcs = common.validateSrcs();
+    List<Artifact> allOutputs = new ArrayList<>(srcs);
+    allOutputs.addAll(semantics.precompiledPythonFiles(ruleContext, srcs, common));
+
+    common.initBinary(allOutputs);
     semantics.validate(ruleContext, common);
     if (ruleContext.hasErrors()) {
       return null;
@@ -127,4 +131,3 @@ public abstract class PyBinary implements RuleConfiguredTargetFactory {
     };
   }
 }
-
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 33f6755041..a537eb936a 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
@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.actions.extra.PythonInfo;
 import com.google.devtools.build.lib.analysis.AnalysisEnvironment;
@@ -31,6 +32,7 @@ import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.Util;
+import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
@@ -341,6 +343,56 @@ public final class PyCommon {
         ruleContext.getPrerequisites("data", Mode.DATA)));
   }
 
+  protected static final ResourceSet PY_COMPILE_RESOURCE_SET =
+      ResourceSet.createWithRamCpuIo(10 /* MB */, 1 /* CPU */, 0.0 /* IO */);
+
+  /**
+   * Utility function to compile multiple .py files to .pyc files.
+   */
+  public Collection<Artifact> createPycFiles(
+      Iterable<Artifact> sources, PathFragment pythonBinary) {
+    List<Artifact> pycFiles = new ArrayList<>();
+    for (Artifact source : sources) {
+      Artifact pycFile = createPycFile(source, pythonBinary);
+      pycFiles.add(pycFile);
+    }
+    return ImmutableList.copyOf(pycFiles);
+  }
+
+  /**
+   * Given a single .py source artifact generate a .pyc file.
+   */
+  public Artifact createPycFile(
+      Artifact source, PathFragment pythonBinary) {
+    Artifact output =
+        ruleContext.getRelatedArtifact(source.getRootRelativePath(), ".pyc");
+
+    // TODO(nnorwitz): Consider adding PYTHONHASHSEED=0 to the environment.
+    // This will make the .pyc more stable, though it will still be non-deterministic.
+    // The timestamp is zeroed out above.
+    SpawnAction.Builder builder = new SpawnAction.Builder()
+        .setResources(PY_COMPILE_RESOURCE_SET)
+        .setExecutable(pythonBinary)
+        .setProgressMessage("Compiling Python")
+        .addInputArgument(
+            ruleContext.getPrerequisiteArtifact("$python_precompile", Mode.HOST))
+        .setMnemonic("PyCompile");
+
+    TransitiveInfoCollection pythonTarget =
+        ruleContext.getPrerequisite(":host_python2_runtime", Mode.HOST);
+    if (pythonTarget != null) {
+      builder.addInputs(pythonTarget
+          .getProvider(FileProvider.class)
+          .getFilesToBuild());
+    }
+
+    builder.addInputArgument(source);
+    builder.addOutputArgument(output);
+    ruleContext.registerAction(builder.build(ruleContext));
+    return output;
+  }
+
+
   /**
    * Returns true if this target has an .so file in its transitive dependency closure.
    */
diff --git a/src/main/java/com/google/devtools/build/lib/rules/python/PyLibrary.java b/src/main/java/com/google/devtools/build/lib/rules/python/PyLibrary.java
index 5f57c1ff6c..fc6cecbaa4 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/python/PyLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/python/PyLibrary.java
@@ -28,6 +28,9 @@ import com.google.devtools.build.lib.rules.cpp.CcLinkParams;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsProvider;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsStore;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * An implementation for the {@code py_library} rule.
  */
@@ -45,8 +48,13 @@ public abstract class PyLibrary implements RuleConfiguredTargetFactory {
     PyCommon common = new PyCommon(ruleContext);
     common.initCommon(common.getDefaultPythonVersion());
     common.validatePackageName();
+
+    List<Artifact> srcs = common.validateSrcs();
+    List<Artifact> allOutputs = new ArrayList<>(srcs);
+    allOutputs.addAll(semantics.precompiledPythonFiles(ruleContext, srcs, common));
+
     NestedSet<Artifact> filesToBuild =
-        NestedSetBuilder.wrap(Order.STABLE_ORDER, common.validateSrcs());
+        NestedSetBuilder.wrap(Order.STABLE_ORDER, allOutputs);
     common.addPyExtraActionPseudoAction();
 
     CcLinkParamsStore ccLinkParamsStore = new CcLinkParamsStore() {
@@ -79,4 +87,3 @@ public abstract class PyLibrary implements RuleConfiguredTargetFactory {
         .build();
   }
 }
-
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 863c3ee76d..357020029b 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
@@ -13,12 +13,15 @@
 // limitations under the License.
 package com.google.devtools.build.lib.rules.python;
 
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.rules.cpp.CcLinkParamsStore;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.InstrumentationSpec;
 
+import java.util.Collection;
+
 /**
  * Pluggable semantics for Python rules.
  *
@@ -47,6 +50,12 @@ public interface PythonSemantics {
   InstrumentationSpec getCoverageInstrumentationSpec();
 
   /**
+   * Utility function to compile multiple .py files to .pyc files, if required.
+   */
+  Collection<Artifact> precompiledPythonFiles(
+      RuleContext ruleContext, Collection<Artifact> sources, PyCommon common);
+
+  /**
    * Create the actual executable artifact.
    *
    * <p>This should create a generating action for {@code common.getExecutable()}.
@@ -59,4 +68,4 @@ public interface PythonSemantics {
    */
   void postInitBinary(RuleContext ruleContext, RunfilesSupport runfilesSupport,
       PyCommon common);
-}
\ No newline at end of file
+}