diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFile.java b/python/psi-api/src/com/jetbrains/python/psi/PyFile.java index 1756436872de..2ec0700aad60 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyFile.java +++ b/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 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. * diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFromImportStatement.java b/python/psi-api/src/com/jetbrains/python/psi/PyFromImportStatement.java index 08a5ea4869ba..d412bfb4293e 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyFromImportStatement.java +++ b/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, diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java b/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java index c9f7ba7553c0..c98621b52d52 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java +++ b/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; diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index a78999c025ac..4467eab49a94 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/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); diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java index 7d9bc1f74f23..31f1a2fd0128 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java +++ b/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 stack = myFindExportedNameStack.get(); diff --git a/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java index ab1d0ffacf43..40110e03e839 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java +++ b/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 candidates = resolveImportSourceCandidates(); - return candidates.size() > 0 ? candidates.get(0) : null; + public PsiFileSystemItem resolveImportSource() { + return as(ContainerUtil.getFirstItem(resolveImportSourceCandidates()), PsiFileSystemItem.class); } @NotNull