Refactor how PyCommon manages its transitive sources

Simplify computing the transitive sources and transitive-without-local sources by factoring a method or two and separating out the computation of `convertedFiles`.

Work toward bazelbuild#7010, yak shaving for bazelbuild#6583 and bazelbuild#1446.

RELNOTES: None
PiperOrigin-RevId: 229410930

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

@@ -56,9 +56,7 @@
 import java.util.Map;
 import java.util.UUID;
 
-/**
- * A helper class for Python rules.
- */
+/** A helper class for analyzing a Python configured target. */
 public final class PyCommon {
 
   public static final String DEFAULT_PYTHON_VERSION_ATTRIBUTE = "default_python_version";
@@ -74,12 +72,6 @@ public void collectMetadataArtifacts(Iterable<Artifact> artifacts,
 
   private final RuleContext ruleContext;
 
-  private Artifact executable = null;
-
-  private final NestedSet<Artifact> transitivePythonSources;
-
-  private final boolean usesSharedLibraries;
-
   /**
    * The Python major version for which this target is being built, as per the {@code
    * python_version} attribute or the configuration.
@@ -94,16 +86,40 @@ public void collectMetadataArtifacts(Iterable<Artifact> artifacts,
    */
   private final PythonVersion sourcesVersion;
 
-  private Map<PathFragment, Artifact> convertedFiles;
+  /**
+   * The Python sources belonging to this target's transitive {@code deps}, not including this
+   * target's own {@code srcs}.
+   */
+  private final NestedSet<Artifact> dependencyTransitivePythonSources;
+
+  /**
+   * The Python sources belonging to this target's transitive {@code deps}, including the Python
+   * sources in this target's {@code srcs}.
+   */
+  private final NestedSet<Artifact> transitivePythonSources;
+
+  /** Whether this target or any of its {@code deps} or {@code data} deps has a shared library. */
+  private final boolean usesSharedLibraries;
+
+  /**
+   * Symlink map from root-relative paths to 2to3 converted source artifacts.
+   *
+   * <p>Null if no 2to3 conversion is required.
+   */
+  private final Map<PathFragment, Artifact> convertedFiles;
+
+  private Artifact executable = null;
 
   private NestedSet<Artifact> filesToBuild = null;
 
   public PyCommon(RuleContext ruleContext) {
     this.ruleContext = ruleContext;
     this.sourcesVersion = getSrcsVersionAttr(ruleContext);
     this.version = ruleContext.getFragment(PythonConfiguration.class).getPythonVersion();
+    this.dependencyTransitivePythonSources = collectDependencyTransitivePythonSources();
     this.transitivePythonSources = collectTransitivePythonSources();
     this.usesSharedLibraries = checkForSharedLibraries();
+    this.convertedFiles = maybeMakeConvertedFiles();
     checkSourceIsCompatible(this.version, this.sourcesVersion, ruleContext.getLabel());
     validatePythonVersionAttr(ruleContext);
   }
@@ -252,28 +268,21 @@ void validatePackageName() {
   }
 
   /**
-   * Adds a {@link PseudoAction} to the build graph that is only used
-   * for providing information to the blaze extra_action feature.
+   * Adds a {@link PseudoAction} to the build graph that is only used for providing information to
+   * the blaze extra_action feature.
    */
   void addPyExtraActionPseudoAction() {
     if (ruleContext.getConfiguration().getActionListeners().isEmpty()) {
       return;
     }
-
-    // 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();
-    collectTransitivePythonSourcesFrom(getTargetDeps(), depBuilder);
-    NestedSet<Artifact> dependencies = depBuilder.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,
+            // We must not add the files declared in the srcs of this rule.;
+            dependencyTransitivePythonSources,
             PseudoAction.getDummyOutput(ruleContext)));
   }
 
@@ -305,16 +314,23 @@ public static Action makePyExtraActionPseudoAction(
   @AutoCodec @AutoCodec.VisibleForSerialization
   static final GeneratedExtension<ExtraActionInfo, PythonInfo> PYTHON_INFO = PythonInfo.pythonInfo;
 
-  private void addSourceFiles(NestedSetBuilder<Artifact> builder, Iterable<Artifact> artifacts) {
-    Preconditions.checkState(convertedFiles == null);
+  /**
+   * If 2to3 conversion is to be done, creates the 2to3 actions and returns the map of converted
+   * files; otherwise returns null.
+   */
+  // TODO(#1393): 2to3 conversion doesn't work in Bazel and the attempt to invoke it for Bazel
+  // should be removed / factored away into PythonSemantics.
+  private Map<PathFragment, Artifact> maybeMakeConvertedFiles() {
     if (sourcesVersion == PythonVersion.PY2 && version == PythonVersion.PY3) {
-      convertedFiles = PythonUtils.generate2to3Actions(ruleContext, artifacts);
+      Iterable<Artifact> artifacts =
+          ruleContext
+              .getPrerequisiteArtifacts("srcs", Mode.TARGET)
+              .filter(PyRuleClasses.PYTHON_SOURCE)
+              .list();
+      return PythonUtils.generate2to3Actions(ruleContext, artifacts);
+    } else {
+      return null;
     }
-    builder.addAll(artifacts);
-  }
-
-  private Iterable<? extends TransitiveInfoCollection> getTargetDeps() {
-    return ruleContext.getPrerequisites("deps", Mode.TARGET);
   }
 
   private static String getOrderErrorMessage(String fieldName, Order expected, Order actual) {
@@ -352,21 +368,34 @@ private void collectTransitivePythonSourcesFrom(
 
   private NestedSet<Artifact> collectTransitivePythonSources() {
     NestedSetBuilder<Artifact> builder = NestedSetBuilder.compileOrder();
-    collectTransitivePythonSourcesFrom(getTargetDeps(), builder);
-    addSourceFiles(builder,
-        ruleContext.getPrerequisiteArtifacts("srcs", Mode.TARGET)
-            .filter(PyRuleClasses.PYTHON_SOURCE).list());
+    collectTransitivePythonSourcesFrom(ruleContext.getPrerequisites("deps", Mode.TARGET), builder);
+    builder.addAll(
+        ruleContext
+            .getPrerequisiteArtifacts("srcs", Mode.TARGET)
+            .filter(PyRuleClasses.PYTHON_SOURCE)
+            .list());
     return builder.build();
   }
 
-  // TODO(brandjon): Move provider merging logic into PyProvider. Have rule implementations read
-  // the sources off a merged provider of deps (with/without the local rule included in the merge).
-  public NestedSet<Artifact> collectTransitivePythonSourcesWithoutLocal() {
+  private NestedSet<Artifact> collectDependencyTransitivePythonSources() {
     NestedSetBuilder<Artifact> builder = NestedSetBuilder.compileOrder();
-    collectTransitivePythonSourcesFrom(getTargetDeps(), builder);
+    collectTransitivePythonSourcesFrom(ruleContext.getPrerequisites("deps", Mode.TARGET), builder);
     return builder.build();
   }
 
+  /** Returns the transitive Python sources collected from the deps and srcs attributes. */
+  public NestedSet<Artifact> getTransitivePythonSources() {
+    return transitivePythonSources;
+  }
+
+  /**
+   * Returns the transitive Python sources collected from the deps attribute, not including sources
+   * from the srcs attribute (unless they were separately reached via deps).
+   */
+  public NestedSet<Artifact> getDependencyTransitivePythonSources() {
+    return dependencyTransitivePythonSources;
+  }
+
   public NestedSet<String> collectImports(RuleContext ruleContext, PythonSemantics semantics) {
     NestedSetBuilder<String> builder = NestedSetBuilder.compileOrder();
     builder.addAll(semantics.getImports(ruleContext));
@@ -375,7 +404,7 @@ public NestedSet<String> collectImports(RuleContext ruleContext, PythonSemantics
   }
 
   private void collectTransitivePythonImports(NestedSetBuilder<String> builder) {
-    for (TransitiveInfoCollection dep : getTargetDeps()) {
+    for (TransitiveInfoCollection dep : ruleContext.getPrerequisites("deps", Mode.TARGET)) {
       if (dep.getProvider(PythonImportsProvider.class) != null) {
         PythonImportsProvider provider = dep.getProvider(PythonImportsProvider.class);
         NestedSet<String> imports = provider.getTransitivePythonImports();