From ad8cef190c57e0b0981fb939f1441653bf5312ea Mon Sep 17 00:00:00 2001 From: Dmitry Jemerov Date: Thu, 6 Mar 2008 19:40:02 +0300 Subject: [PATCH] better resolve for imported elements in python --- .../python/psi/PyImportStatement.java | 2 +- .../jetbrains/python/psi/PyResolveUtil.java | 10 +++++++- .../psi/impl/PyFromImportStatementImpl.java | 2 +- .../python/psi/impl/PyImportElementImpl.java | 7 +++++- .../psi/impl/PyImportStatementImpl.java | 23 +++++++++++------- .../psi/impl/PyReferenceExpressionImpl.java | 8 +++---- .../python/psi/impl/ResolveImportUtil.java | 24 +++++++++++++++---- 7 files changed, 55 insertions(+), 21 deletions(-) diff --git a/python/src/com/jetbrains/python/psi/PyImportStatement.java b/python/src/com/jetbrains/python/psi/PyImportStatement.java index ee1572f59645..da795179646f 100644 --- a/python/src/com/jetbrains/python/psi/PyImportStatement.java +++ b/python/src/com/jetbrains/python/psi/PyImportStatement.java @@ -24,5 +24,5 @@ package com.jetbrains.python.psi; * To change this template use File | Settings | File Templates. */ public interface PyImportStatement extends PyStatement { - PyImportElement getImportElement(); + PyImportElement[] getImportElements(); } diff --git a/python/src/com/jetbrains/python/psi/PyResolveUtil.java b/python/src/com/jetbrains/python/psi/PyResolveUtil.java index 92fc5e068d31..ef654767a29f 100644 --- a/python/src/com/jetbrains/python/psi/PyResolveUtil.java +++ b/python/src/com/jetbrains/python/psi/PyResolveUtil.java @@ -20,6 +20,7 @@ import com.intellij.codeInsight.lookup.LookupElement; import com.intellij.codeInsight.lookup.LookupElementFactory; import com.intellij.psi.*; import com.intellij.psi.scope.PsiScopeProcessor; +import com.intellij.openapi.vfs.VirtualFile; import org.jetbrains.annotations.Nullable; import java.util.ArrayList; @@ -70,7 +71,14 @@ public class PyResolveUtil { } public boolean execute(PsiElement element, ResolveState substitutor) { - if (element instanceof PsiNamedElement) { + if (element instanceof PyFile) { + final VirtualFile file = ((PyFile)element).getVirtualFile(); + if (file != null && _name.equals(file.getNameWithoutExtension())) { + _result = element; + return false; + } + } + else if (element instanceof PsiNamedElement) { if (_name.equals(((PsiNamedElement)element).getName())) { _result = element; return false; diff --git a/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java index 5586a9f03e2d..c96d49ccf982 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java @@ -77,7 +77,7 @@ public class PyFromImportStatementImpl extends PyElementImpl implements PyFromIm if (isStarImport()) { PyReferenceExpression expr = getImportSource(); if (expr != null) { - final PsiElement importedFile = ResolveImportUtil.resolveImportReference(expr, expr.getReferencedName()); + final PsiElement importedFile = ResolveImportUtil.resolveImportReference(expr); if (importedFile != null) { return importedFile.processDeclarations(processor, state, null, place); } diff --git a/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java index 824bcee3134b..4ee20ce49339 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java @@ -44,6 +44,11 @@ public class PyImportElementImpl extends PyElementImpl implements PyImportElemen @Override public boolean processDeclarations(@NotNull final PsiScopeProcessor processor, @NotNull final ResolveState state, final PsiElement lastParent, @NotNull final PsiElement place) { - return processor.execute(getImportReference(), state); + final PyReferenceExpression importRef = getImportReference(); + final PsiElement element = importRef.resolve(); + if (element != null) { + return processor.execute(element, state); + } + return true; } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyImportStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyImportStatementImpl.java index b2c7bf9a189d..da2c5798d408 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyImportStatementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyImportStatementImpl.java @@ -20,10 +20,11 @@ import com.intellij.lang.ASTNode; import com.intellij.psi.PsiElement; import com.intellij.psi.ResolveState; import com.intellij.psi.scope.PsiScopeProcessor; -import org.jetbrains.annotations.NotNull; +import com.intellij.psi.tree.TokenSet; import com.jetbrains.python.PyElementTypes; import com.jetbrains.python.psi.PyImportElement; import com.jetbrains.python.psi.PyImportStatement; +import org.jetbrains.annotations.NotNull; /** * Created by IntelliJ IDEA. @@ -33,17 +34,23 @@ import com.jetbrains.python.psi.PyImportStatement; * To change this template use File | Settings | File Templates. */ public class PyImportStatementImpl extends PyElementImpl implements PyImportStatement { - public PyImportStatementImpl(ASTNode astNode) { - super(astNode); - } + public PyImportStatementImpl(ASTNode astNode) { + super(astNode); + } @Override - public boolean processDeclarations(@NotNull final PsiScopeProcessor processor, @NotNull final ResolveState state, final PsiElement lastParent, + public boolean processDeclarations(@NotNull final PsiScopeProcessor processor, + @NotNull final ResolveState state, + final PsiElement lastParent, @NotNull final PsiElement place) { - return getImportElement().processDeclarations(processor, state, lastParent, place); + for (PyImportElement element : getImportElements()) { + if (element == lastParent) continue; + if (!element.processDeclarations(processor, state, null, place)) return false; + } + return true; } - public PyImportElement getImportElement() { - return (PyImportElement)getNode().findChildByType(PyElementTypes.IMPORT_ELEMENT).getPsi(); + public PyImportElement[] getImportElements() { + return childrenToPsi(TokenSet.create(PyElementTypes.IMPORT_ELEMENT), new PyImportElement[0]); } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java index b9d494b2df98..8f685e081ab8 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java @@ -96,6 +96,10 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere final String referencedName = getReferencedName(); if (referencedName == null) return null; + if (getParent() instanceof PyImportElement || getParent() instanceof PyFromImportStatement) { + return ResolveImportUtil.resolveImportReference(this); + } + final PyExpression qualifier = getQualifier(); if (qualifier != null) { if (qualifier instanceof PyCallExpression) { @@ -125,10 +129,6 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere return null; // TODO? } - if (getParent() instanceof PyImportElement || getParent() instanceof PyFromImportStatement) { - return ResolveImportUtil.resolveImportReference(this, referencedName); - } - return PyResolveUtil.treeWalkUp(new PyResolveUtil.ResolveProcessor(referencedName), this, this, null); } diff --git a/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java b/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java index 1cee9af814a4..bf15caf9f5f7 100644 --- a/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java +++ b/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java @@ -1,11 +1,11 @@ package com.jetbrains.python.psi.impl; -import org.jetbrains.annotations.Nullable; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; import com.intellij.psi.search.FilenameIndex; import com.intellij.psi.search.GlobalSearchScope; -import com.jetbrains.python.psi.PyElement; +import com.jetbrains.python.psi.*; +import org.jetbrains.annotations.Nullable; /** * @author yole @@ -15,9 +15,23 @@ public class ResolveImportUtil { } @Nullable - static PsiElement resolveImportReference(final PyElement context, final String referencedName) { - final PsiFile[] files = FilenameIndex.getFilesByName(context.getProject(), referencedName + ".py", - GlobalSearchScope.allScope(context.getProject())); + static PsiElement resolveImportReference(final PyReferenceExpression importRef) { + String referencedName = importRef.getReferencedName(); + if (referencedName == null) return null; + if (importRef.getParent() instanceof PyImportElement) { + PyImportElement parent = (PyImportElement) importRef.getParent(); + if (parent.getParent() instanceof PyFromImportStatement) { + PyFromImportStatement stmt = (PyFromImportStatement) parent.getParent(); + final PyReferenceExpression source = stmt.getImportSource(); + if (source == null) return null; + PsiElement sourceFile = resolveImportReference(source); + if (sourceFile instanceof PyFile) { + return PyResolveUtil.treeWalkUp(new PyResolveUtil.ResolveProcessor(referencedName), sourceFile, null, importRef); + } + } + } + final PsiFile[] files = FilenameIndex.getFilesByName(importRef.getProject(), referencedName + ".py", + GlobalSearchScope.allScope(importRef.getProject())); if (files.length == 1) return files[0]; return null; }