Several improvements in Python PSI API related to imports

* Pull PyFileImpl#findExportedName up to PyFile interface
* PyFromImportStatement#resolveImportSource returns PsiFileSystemItem

python/psi-api/src/com/jetbrains/python/psi/PyFile.java

@@ -15,6 +15,7 @@
  */
 package com.jetbrains.python.psi;
 
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import org.jetbrains.annotations.NotNull;
@@ -51,6 +52,9 @@ public interface PyFile extends PyElement, PsiFile, PyDocStringOwner, ScopeOwner
   @NotNull
   List<PyFromImportStatement> getFromImports();
 
+  @Nullable
+  PsiElement findExportedName(String name);
+
   /**
    * Returns the list of import elements in all 'import xxx' statements in the top-level scope of the file.
    *

python/psi-api/src/com/jetbrains/python/psi/PyFromImportStatement.java

@@ -16,6 +16,7 @@
 package com.jetbrains.python.psi;
 
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFileSystemItem;
 import com.intellij.psi.StubBasedPsiElement;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.stubs.PyFromImportStatementStub;
@@ -77,7 +78,7 @@ public interface PyFromImportStatement extends PyImportStatementBase, StubBasedP
    * @return the resolved import source (file or directory containing __init__.py), or null if the import is unresolved.
    */
   @Nullable
-  PsiElement resolveImportSource();
+  PsiFileSystemItem resolveImportSource();
 
   /**
    * Resolves the import source qualified name to a number of possible files or directories. Note: performs a Python only resolve,

python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java

@@ -179,7 +179,7 @@ public class PythonReferenceImporter implements ReferenceImporter {
         existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, importElement, importElement.resolve());
       }
       for (PyFromImportStatement fromImportStatement : pyFile.getFromImports()) {
-        if (!(fromImportStatement.isStarImport()) && fromImportStatement.getImportElements().length > 0) {
+        if (!fromImportStatement.isStarImport() && fromImportStatement.getImportElements().length > 0) {
           PsiElement source = fromImportStatement.resolveImportSource();
           existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, fromImportStatement.getImportElements()[0], source);
         }
@@ -194,7 +194,7 @@ public class PythonReferenceImporter implements ReferenceImporter {
                                              PsiFile existingImportFile,
                                              PyImportElement importElement,
                                              PsiElement source) {
-    PsiElement sourceFile = PyUtil.turnDirIntoInit(source);
+    PyFile sourceFile = ((PyFile)PyUtil.turnDirIntoInit(source));
     if (sourceFile instanceof PyFileImpl) {
       PyStatement importStatement = importElement.getContainingImportStatement();
       String refName = null;
@@ -217,13 +217,12 @@ public class PythonReferenceImporter implements ReferenceImporter {
         seenFileNames.add(refName);
       }
 
-      PyFileImpl importSourceFile = (PyFileImpl)sourceFile;
-      PsiElement res = importSourceFile.findExportedName(refText);
+      PsiElement res = sourceFile.findExportedName(refText);
       // allow importing from this source if it either declares the name itself or represents a higher-level package that reexports the name
       if (res != null && !(res instanceof PyFile) && !(res instanceof PyImportElement) && res.getContainingFile() != null &&
           PsiTreeUtil.isAncestor(source, res.getContainingFile(), false)) {
-        existingImportFile = importSourceFile;
-        fix.addImport(res, importSourceFile, importElement);
+        existingImportFile = sourceFile;
+        fix.addImport(res, sourceFile, importElement);
       }
     }
     return existingImportFile;

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -889,7 +889,7 @@ public class PyUtil {
    * @return a PsiFile if target was a PsiDirectory, or null, or target unchanged.
    */
   @Nullable
-  public static PsiElement turnDirIntoInit(PsiElement target) {
+  public static PsiElement turnDirIntoInit(@Nullable PsiElement target) {
     if (target instanceof PsiDirectory) {
       final PsiDirectory dir = (PsiDirectory)target;
       final PsiFile file = dir.findFile(PyNames.INIT_DOT_PY);

python/src/com/jetbrains/python/psi/impl/PyFileImpl.java

@@ -430,6 +430,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
     return PyPsiUtils.collectStubChildren(this, this.getStub(), PyElementTypes.TARGET_EXPRESSION, PyTargetExpression.class);
   }
 
+  @Override
   @Nullable
   public PsiElement findExportedName(String name) {
     final List<String> stack = myFindExportedNameStack.get();

python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java

@@ -22,6 +22,7 @@ import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayFactory;
 import com.intellij.util.ArrayUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
@@ -36,6 +37,8 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
+import static com.jetbrains.python.psi.PyUtil.as;
+
 /**
  * @author yole
  */
@@ -216,9 +219,8 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
   }
 
   @Nullable
-  public PsiElement resolveImportSource() {
-    final List<PsiElement> candidates = resolveImportSourceCandidates();
-    return candidates.size() > 0 ? candidates.get(0) : null;
+  public PsiFileSystemItem resolveImportSource() {
+    return as(ContainerUtil.getFirstItem(resolveImportSourceCandidates()), PsiFileSystemItem.class);
   }
 
   @NotNull