1 files changed, 42 insertions, 17 deletions

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 ed892d3ff0..bf67b4f7aa 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
@@ -17,6 +17,7 @@ import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ResourceSet;
@@ -203,27 +204,46 @@ public final class PyCommon {
       return;
     }
 
-    // Has to be unfiltered sources as filtered will give an error for
-    // unsupported file types where as certain tests only expect a warning.
-    Collection<Artifact> sources = ruleContext.getPrerequisiteArtifacts("srcs", Mode.TARGET).list();
-
     // We need to do it in this convoluted way because we must not add the files declared in the
     // srcs of this rule. Note that it is not enough to remove the direct members from the nested
     // set of the current rule, because the same files may have been declared in a dependency, too.
     NestedSetBuilder<Artifact> depBuilder = NestedSetBuilder.compileOrder();
-    collectTransitivePythonSourcesFromDeps(depBuilder);
+    collectTransitivePythonSourcesFrom(getTargetDeps(), depBuilder);
     NestedSet<Artifact> dependencies = depBuilder.build();
 
-    PythonInfo info = PythonInfo.newBuilder()
-        .addAllSourceFile(Artifact.toExecPaths(sources))
-        .addAllDepFile(Artifact.toExecPaths(dependencies))
-        .build();
+    ruleContext.registerAction(
+        makePyExtraActionPseudoAction(
+            ruleContext.getActionOwner(),
+            // Has to be unfiltered sources as filtered will give an error for
+            // unsupported file types where as certain tests only expect a warning.
+            ruleContext.getPrerequisiteArtifacts("srcs", Mode.TARGET).list(),
+            dependencies,
+            PseudoAction.getDummyOutput(ruleContext)));
+  }
 
-    ruleContext.getAnalysisEnvironment()
-        .registerAction(new PyPseudoAction(ruleContext.getActionOwner(),
-            NestedSetBuilder.wrap(Order.STABLE_ORDER, Iterables.concat(sources, dependencies)),
-            ImmutableList.of(PseudoAction.getDummyOutput(ruleContext)), "Python",
-            PythonInfo.pythonInfo, info));
+  /**
+   * Creates a {@link PseudoAction} that is only used for providing
+   * information to the blaze extra_action feature.
+   */
+  public static Action makePyExtraActionPseudoAction(
+      ActionOwner owner,
+      Iterable<Artifact> sources,
+      Iterable<Artifact> dependencies,
+      Artifact output) {
+
+    PythonInfo info =
+        PythonInfo.newBuilder()
+            .addAllSourceFile(Artifact.toExecPaths(sources))
+            .addAllDepFile(Artifact.toExecPaths(dependencies))
+            .build();
+
+    return new PyPseudoAction(
+        owner,
+        NestedSetBuilder.wrap(Order.STABLE_ORDER, Iterables.concat(sources, dependencies)),
+        ImmutableList.of(output),
+        "Python",
+        PythonInfo.pythonInfo,
+        info);
   }
 
   private void addSourceFiles(NestedSetBuilder<Artifact> builder, Iterable<Artifact> artifacts) {
@@ -234,8 +254,13 @@ public final class PyCommon {
     builder.addAll(artifacts);
   }
 
-  private void collectTransitivePythonSourcesFromDeps(NestedSetBuilder<Artifact> builder) {
-    for (TransitiveInfoCollection dep : ruleContext.getPrerequisites("deps", Mode.TARGET)) {
+  private Iterable<? extends TransitiveInfoCollection> getTargetDeps() {
+    return ruleContext.getPrerequisites("deps", Mode.TARGET);
+  }
+
+  private void collectTransitivePythonSourcesFrom(
+      Iterable<? extends TransitiveInfoCollection> deps, NestedSetBuilder<Artifact> builder) {
+    for (TransitiveInfoCollection dep : deps) {
       if (dep.getProvider(PythonSourcesProvider.class) != null) {
         PythonSourcesProvider provider = dep.getProvider(PythonSourcesProvider.class);
         builder.addTransitive(provider.getTransitivePythonSources());
@@ -251,7 +276,7 @@ public final class PyCommon {
   private NestedSet<Artifact> collectTransitivePythonSources() {
     NestedSetBuilder<Artifact> builder =
         NestedSetBuilder.compileOrder();
-    collectTransitivePythonSourcesFromDeps(builder);
+    collectTransitivePythonSourcesFrom(getTargetDeps(), builder);
     addSourceFiles(builder, ruleContext
         .getPrerequisiteArtifacts("srcs", Mode.TARGET).filter(PyRuleClasses.PYTHON_SOURCE).list());
     return builder.build();