1 files changed, 21 insertions, 4 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 969955f55a..9b6813fad5 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
@@ -276,14 +276,31 @@ public final class PyCommon {
   }
 
   private NestedSet<Artifact> collectTransitivePythonSources() {
-    NestedSetBuilder<Artifact> builder =
-        NestedSetBuilder.compileOrder();
+    NestedSetBuilder<Artifact> builder = NestedSetBuilder.compileOrder();
     collectTransitivePythonSourcesFrom(getTargetDeps(), builder);
-    addSourceFiles(builder, ruleContext
-        .getPrerequisiteArtifacts("srcs", Mode.TARGET).filter(PyRuleClasses.PYTHON_SOURCE).list());
+    addSourceFiles(builder,
+        ruleContext.getPrerequisiteArtifacts("srcs", Mode.TARGET)
+            .filter(PyRuleClasses.PYTHON_SOURCE).list());
     return builder.build();
   }
 
+  public NestedSet<PathFragment> collectImports(
+      RuleContext ruleContext, PythonSemantics semantics) {
+    NestedSetBuilder<PathFragment> builder = NestedSetBuilder.compileOrder();
+    builder.addAll(semantics.getImports(ruleContext));
+    collectTransitivePythonImports(builder);
+    return builder.build();
+  }
+
+  private void collectTransitivePythonImports(NestedSetBuilder<PathFragment> builder) {
+    for (TransitiveInfoCollection dep : getTargetDeps()) {
+      if (dep.getProvider(PythonImportsProvider.class) != null) {
+        PythonImportsProvider provider = dep.getProvider(PythonImportsProvider.class);
+        builder.addTransitive(provider.getTransitivePythonImports());
+      }
+    }
+  }
+
   /**
    * Checks that the source file version is compatible with the Python interpreter.
    */