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 371f83a08204..28c3c2083aa0 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyFile.java +++ b/python/psi-api/src/com/jetbrains/python/psi/PyFile.java @@ -27,10 +27,16 @@ public interface PyFile extends PyElement, PsiFile, PyDocStringOwner, ScopeOwner LanguageLevel getLanguageLevel(); + /** + * Return the list of all 'from ... import' statements in the top-level scope of the file. + * + * @return the list of 'from ... import' statements. + */ + @NotNull List getFromImports(); /** - * Returns the list of import elements in all 'import xxx' statements within the file. + * Returns the list of import elements in all 'import xxx' statements in the top-level scope of the file. * * @return the list of import targets. */ @@ -51,7 +57,7 @@ public interface PyFile extends PyElement, PsiFile, PyDocStringOwner, ScopeOwner boolean hasImportFromFuture(FutureFeature feature); /** - * If the function raises a DeprecationWarning or a PendingDeprecationWarning, returns the explanation text provided for the warning.. + * If the function raises a DeprecationWarning or a PendingDeprecationWarning, returns the explanation text provided for the warning. * * @return the deprecation message or null if the function is not deprecated. */ diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java b/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java index 13ccbe8c77ff..8baf23991d82 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PythonReferenceImporter.java @@ -17,9 +17,7 @@ import com.intellij.openapi.util.io.FileUtil; import com.intellij.psi.*; import com.intellij.psi.search.FilenameIndex; import com.intellij.psi.search.GlobalSearchScope; -import com.intellij.psi.tree.TokenSet; import com.intellij.psi.util.PsiTreeUtil; -import com.jetbrains.python.PyElementTypes; import com.jetbrains.python.PyNames; import com.jetbrains.python.codeInsight.PyCodeInsightSettings; import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache; @@ -27,8 +25,6 @@ import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyFileImpl; import com.jetbrains.python.psi.impl.PyQualifiedName; -import com.jetbrains.python.psi.resolve.CollectProcessor; -import com.jetbrains.python.psi.resolve.PyResolveUtil; import com.jetbrains.python.psi.resolve.QualifiedNameFinder; import com.jetbrains.python.psi.search.PyProjectScopeBuilder; import com.jetbrains.python.psi.stubs.PyClassNameIndex; @@ -43,7 +39,9 @@ import java.util.*; public class PythonReferenceImporter implements ReferenceImporter { @Override public boolean autoImportReferenceAtCursor(@NotNull final Editor editor, @NotNull final PsiFile file) { - if (!(file instanceof PyFile)) return false; + if (!(file instanceof PyFile)) { + return false; + } int caretOffset = editor.getCaretModel().getOffset(); Document document = editor.getDocument(); int lineNumber = document.getLineNumber(caretOffset); @@ -71,7 +69,9 @@ public class PythonReferenceImporter implements ReferenceImporter { @Override public boolean autoImportReferenceAt(@NotNull Editor editor, @NotNull PsiFile file, int offset) { - if (!(file instanceof PyFile)) return false; + if (!(file instanceof PyFile)) { + return false; + } PsiReference element = file.findReferenceAt(offset); if (element instanceof PyReferenceExpression && isImportable((PsiElement)element)) { final PyReferenceExpression refExpr = (PyReferenceExpression)element; @@ -89,8 +89,6 @@ public class PythonReferenceImporter implements ReferenceImporter { return false; } - private static final TokenSet IS_IMPORT_STATEMENT = TokenSet.create(PyElementTypes.IMPORT_STATEMENT); - @Nullable public static AutoImportQuickFix proposeImportFix(final PyElement node, PsiReference reference) { final String text = reference.getElement().getText(); @@ -109,37 +107,15 @@ public class PythonReferenceImporter implements ReferenceImporter { } AutoImportQuickFix fix = new AutoImportQuickFix(node, reference, !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT); - Set seen_file_names = new HashSet(); // true import names - // maybe the name is importable via some existing 'import foo' statement, and only needs a qualifier. - // walk up collecting all such statements and analyzing - CollectProcessor import_prc = new CollectProcessor(IS_IMPORT_STATEMENT); - PyResolveUtil.treeCrawlUp(import_prc, node); - List result = import_prc.getResult(); - PsiFile existing_import_file = null; // if there's a matching existing import, this it the file it imports - if (!result.isEmpty()) { - for (PsiElement stmt : import_prc.getResult()) { - for (PyImportElement ielt : ((PyImportStatement)stmt).getImportElements()) { - final PyReferenceExpression src = ielt.getImportReferenceExpression(); - if (src != null) { - PsiElement dst = src.getReference().resolve(); - if (dst instanceof PyFileImpl) { - PyFileImpl dstFile = (PyFileImpl)dst; - String name = ielt.getImportReferenceExpression().getReferencedName(); // ref is ok or matching would fail - seen_file_names.add(name); - PsiElement res = dstFile.findExportedName(refText); - if (res != null && !(res instanceof PyFile) && !(res instanceof PyImportElement) && dstFile.equals(res.getContainingFile())) { - existing_import_file = dstFile; - fix.addImport(res, dstFile, ielt); - } - } - } - } - } + Set seenFileNames = new HashSet(); // true import names + + PsiFile existingImportFile = addCandidatesFromExistingImports(node, refText, fix, seenFileNames); + if (fix.getCandidatesCount() == 0) { + // maybe some unimported file has it, too + ProgressManager.checkCanceled(); // before expensive index searches + addSymbolImportCandidates(node, refText, fix, seenFileNames, existingImportFile); } - // maybe some unimported file has it, too - ProgressManager.checkCanceled(); // before expensive index searches - addSymbolImportCandidates(node, refText, fix, seen_file_names, existing_import_file); for(PyImportCandidateProvider provider: Extensions.getExtensions(PyImportCandidateProvider.EP_NAME)) { provider.addImportCandidates(reference, refText, fix); } @@ -150,6 +126,59 @@ public class PythonReferenceImporter implements ReferenceImporter { return null; } + /** + * maybe the name is importable via some existing 'import foo' statement, and only needs a qualifier. + * collect all such statements and analyze. + * NOTE: It only makes sense to look at imports in file scope - there is no guarantee that an import in a local scope will + * be visible from the scope where the auto-import was invoked + * + * @param node + * @param refText + * @param fix + * @param seenFileNames + * @return + */ + @Nullable + private static PsiFile addCandidatesFromExistingImports(PyElement node, String refText, AutoImportQuickFix fix, + Set seenFileNames) { + PsiFile existingImportFile = null; // if there's a matching existing import, this it the file it imports + PsiFile file = node.getContainingFile(); + if (file instanceof PyFile) { + PyFile pyFile = (PyFile)file; + for (PyImportElement importElement : pyFile.getImportTargets()) { + existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, importElement, importElement.resolve()); + } + for (PyFromImportStatement fromImportStatement : pyFile.getFromImports()) { + if (!(fromImportStatement.isStarImport()) && fromImportStatement.getImportElements().length > 0) { + PsiElement source = fromImportStatement.resolveImportSource(); + existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, fromImportStatement.getImportElements()[0], source); + } + } + } + return existingImportFile; + } + + private static PsiFile addImportViaElement(String refText, + AutoImportQuickFix fix, + Set seenFileNames, + PsiFile existingImportFile, + PyImportElement importElement, + PsiElement source) { + PsiElement sourceFile = PyUtil.turnDirIntoInit(source); + if (sourceFile instanceof PyFileImpl) { + seenFileNames.add(importElement.getImportReferenceExpression().getReferencedName()); + PyFileImpl importSourceFile = (PyFileImpl)sourceFile; + PsiElement res = importSourceFile.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) && + PsiTreeUtil.isAncestor(source, res.getContainingFile(), false)) { + existingImportFile = importSourceFile; + fix.addImport(res, importSourceFile, importElement); + } + } + return existingImportFile; + } + private static void addSymbolImportCandidates(PyElement node, String refText, AutoImportQuickFix fix, diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java index 0794b7eda772..6355781b5db9 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java @@ -523,6 +523,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { return ret; } + @NotNull public List getFromImports() { return PyPsiUtils.collectStubChildren(this, getStub(), PyElementTypes.FROM_IMPORT_STATEMENT, PyFromImportStatement.class); } diff --git a/python/src/com/jetbrains/python/psi/resolve/CollectProcessor.java b/python/src/com/jetbrains/python/psi/resolve/CollectProcessor.java deleted file mode 100644 index 8ff69dfcad6d..000000000000 --- a/python/src/com/jetbrains/python/psi/resolve/CollectProcessor.java +++ /dev/null @@ -1,44 +0,0 @@ -package com.jetbrains.python.psi.resolve; - -import com.intellij.openapi.util.Key; -import com.intellij.psi.PsiElement; -import com.intellij.psi.ResolveState; -import com.intellij.psi.tree.TokenSet; -import org.jetbrains.annotations.NotNull; - -import java.util.ArrayList; -import java.util.List; - -public class CollectProcessor implements PyClassScopeProcessor { - private final List myResult; - private final TokenSet myTargetTokenSet; - - public CollectProcessor(TokenSet targetTokenSet) { - myTargetTokenSet = targetTokenSet; - myResult = new ArrayList(); - } - - public boolean execute(@NotNull final PsiElement element, final ResolveState state) { - if (myTargetTokenSet.contains(element.getNode().getElementType())) { - myResult.add(element); - } - return true; // collect till we drop - } - - public T getHint(@NotNull final Key hintKey) { - return null; - } - - public void handleEvent(final Event event, final Object associated) { - } - - public List getResult() { - return myResult; - } - - @NotNull - @Override - public TokenSet getTargetTokenSet() { - return myTargetTokenSet; - } -} diff --git a/python/src/com/jetbrains/python/psi/resolve/PyClassScopeProcessor.java b/python/src/com/jetbrains/python/psi/resolve/PyClassScopeProcessor.java deleted file mode 100644 index 73e80b941b61..000000000000 --- a/python/src/com/jetbrains/python/psi/resolve/PyClassScopeProcessor.java +++ /dev/null @@ -1,18 +0,0 @@ -package com.jetbrains.python.psi.resolve; - -import com.intellij.psi.scope.PsiScopeProcessor; -import com.intellij.psi.tree.TokenSet; -import org.jetbrains.annotations.NotNull; - -/** - * Processor capable of giving multiple hints on what it's looking for. - * User: dcheryasov - * Date: Apr 19, 2009 - */ -public interface PyClassScopeProcessor extends PsiScopeProcessor { - /** - * @return set of element types that might be interesting for the processor. - */ - @NotNull - TokenSet getTargetTokenSet(); -} diff --git a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java b/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java index bd696ee2751e..de0fa7f9e216 100644 --- a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java +++ b/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java @@ -73,14 +73,9 @@ public class PyResolveUtil { } @Nullable - public static PsiElement getPrevNodeOf(PsiElement elt, PsiScopeProcessor proc) { + public static PsiElement getPrevNodeOf(PsiElement elt) { if (elt instanceof PsiFile) return null; // no sense to get the previous node of a file - if (proc instanceof PyClassScopeProcessor) { - return getPrevNodeOf(elt, ((PyClassScopeProcessor)proc).getTargetTokenSet()); - } - else { - return getPrevNodeOf(elt, PythonDialectsTokenSetProvider.INSTANCE.getNameDefinerTokens()); - } + return getPrevNodeOf(elt, PythonDialectsTokenSetProvider.INSTANCE.getNameDefinerTokens()); } /** @@ -163,10 +158,10 @@ public class PyResolveUtil { final boolean is_outside_param_list = PsiTreeUtil.getParentOfType(elt, PyParameterList.class) == null; do { ProgressManager.checkCanceled(); - seeker = getPrevNodeOf(seeker, processor); + seeker = getPrevNodeOf(seeker); // aren't we in the same defining assignment, global, etc? if ((seeker instanceof NameDefiner) && ((NameDefiner)seeker).mustResolveOutside() && PsiTreeUtil.isAncestor(seeker, elt, true)) { - seeker = getPrevNodeOf(seeker, processor); + seeker = getPrevNodeOf(seeker); } // maybe we're under a cap? while (true) { @@ -182,7 +177,7 @@ public class PyResolveUtil { seeker = local_cap; } else { - seeker = getPrevNodeOf(local_cap, processor); + seeker = getPrevNodeOf(local_cap); } } else {