diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java index 8f685e081ab8..4a50c3d8143a 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java @@ -96,7 +96,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere final String referencedName = getReferencedName(); if (referencedName == null) return null; - if (getParent() instanceof PyImportElement || getParent() instanceof PyFromImportStatement) { + if (PsiTreeUtil.getParentOfType(this, PyImportElement.class, PyFromImportStatement.class) != null) { return ResolveImportUtil.resolveImportReference(this); } diff --git a/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java b/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java index bf15caf9f5f7..98109fe9daa5 100644 --- a/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java +++ b/python/src/com/jetbrains/python/psi/impl/ResolveImportUtil.java @@ -1,7 +1,13 @@ package com.jetbrains.python.psi.impl; +import com.intellij.openapi.module.Module; +import com.intellij.openapi.module.ModuleUtil; +import com.intellij.openapi.roots.ModuleRootManager; +import com.intellij.openapi.vfs.VirtualFile; +import com.intellij.psi.PsiDirectory; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; +import com.intellij.psi.PsiManager; import com.intellij.psi.search.FilenameIndex; import com.intellij.psi.search.GlobalSearchScope; import com.jetbrains.python.psi.*; @@ -18,6 +24,14 @@ public class ResolveImportUtil { static PsiElement resolveImportReference(final PyReferenceExpression importRef) { String referencedName = importRef.getReferencedName(); if (referencedName == null) return null; + + final PyExpression qualifier = importRef.getQualifier(); + if (qualifier instanceof PyReferenceExpression) { + PsiElement qualifierElement = ((PyReferenceExpression) qualifier).resolve(); + if (qualifierElement == null) return null; + return resolveChild(qualifierElement, referencedName, importRef); + } + if (importRef.getParent() instanceof PyImportElement) { PyImportElement parent = (PyImportElement) importRef.getParent(); if (parent.getParent() instanceof PyFromImportStatement) { @@ -25,14 +39,40 @@ public class ResolveImportUtil { 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); - } + return resolveChild(sourceFile, referencedName, importRef); } } final PsiFile[] files = FilenameIndex.getFilesByName(importRef.getProject(), referencedName + ".py", GlobalSearchScope.allScope(importRef.getProject())); if (files.length == 1) return files[0]; + + final Module module = ModuleUtil.findModuleForPsiElement(importRef); + final VirtualFile[] contentRoots = ModuleRootManager.getInstance(module).getContentRoots(); + for(VirtualFile root: contentRoots) { + final VirtualFile childFile = root.findChild(referencedName + ".py"); + if (childFile != null) { + return PsiManager.getInstance(importRef.getProject()).findFile(childFile); + } + + final VirtualFile childDir = root.findChild(referencedName); + if (childDir != null) { + return PsiManager.getInstance(importRef.getProject()).findDirectory(childDir); + } + } + return null; } + + @Nullable + private static PsiElement resolveChild(final PsiElement parent, final String referencedName, final PyReferenceExpression importRef) { + if (parent instanceof PyFile) { + return PyResolveUtil.treeWalkUp(new PyResolveUtil.ResolveProcessor(referencedName), parent, null, importRef); + } + else if (parent instanceof PsiDirectory) { + return ((PsiDirectory)parent).findFile(referencedName + ".py"); + } + else { + return null; + } + } } diff --git a/python/testData/resolve/multiFile/fromPackageImport/FromPackageImport.py b/python/testData/resolve/multiFile/fromPackageImport/FromPackageImport.py new file mode 100644 index 000000000000..4776ec01845e --- /dev/null +++ b/python/testData/resolve/multiFile/fromPackageImport/FromPackageImport.py @@ -0,0 +1,2 @@ +from mypackage import p +# \ No newline at end of file diff --git a/python/testData/resolve/multiFile/fromPackageImport/mypackage/__init__.py b/python/testData/resolve/multiFile/fromPackageImport/mypackage/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromPackageImport/mypackage/p.py b/python/testData/resolve/multiFile/fromPackageImport/mypackage/p.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromPackageImportFile/FromPackageImportFile.py b/python/testData/resolve/multiFile/fromPackageImportFile/FromPackageImportFile.py new file mode 100644 index 000000000000..faafe5f7d809 --- /dev/null +++ b/python/testData/resolve/multiFile/fromPackageImportFile/FromPackageImportFile.py @@ -0,0 +1,2 @@ +from mypackage import myfile +# \ No newline at end of file diff --git a/python/testData/resolve/multiFile/fromPackageImportFile/mypackage/__init__.py b/python/testData/resolve/multiFile/fromPackageImportFile/mypackage/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromPackageImportFile/mypackage/myfile.py b/python/testData/resolve/multiFile/fromPackageImportFile/mypackage/myfile.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImport/FromQualifiedPackageImport.py b/python/testData/resolve/multiFile/fromQualifiedPackageImport/FromQualifiedPackageImport.py new file mode 100644 index 000000000000..dd933343fae3 --- /dev/null +++ b/python/testData/resolve/multiFile/fromQualifiedPackageImport/FromQualifiedPackageImport.py @@ -0,0 +1,2 @@ +from mypackage.myfile import myclass +# \ No newline at end of file diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/__init__.py b/python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/myfile.py b/python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/myfile.py new file mode 100644 index 000000000000..45b107f8475a --- /dev/null +++ b/python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/myfile.py @@ -0,0 +1,2 @@ +class myclass: + pass \ No newline at end of file diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/FromQualifiedPackageImportFile.py b/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/FromQualifiedPackageImportFile.py new file mode 100644 index 000000000000..470c830aa62b --- /dev/null +++ b/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/FromQualifiedPackageImportFile.py @@ -0,0 +1,2 @@ +from mypackage.myfile import myclass +# \ No newline at end of file diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/__init__.py b/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/myfile.py b/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/myfile.py new file mode 100644 index 000000000000..45b107f8475a --- /dev/null +++ b/python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/myfile.py @@ -0,0 +1,2 @@ +class myclass: + pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java index 9f50245b3854..c7849c7908fc 100644 --- a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java @@ -4,10 +4,7 @@ import com.intellij.codeInsight.CodeInsightTestCase; import com.intellij.openapi.application.PathManager; import com.intellij.openapi.editor.Document; import com.intellij.openapi.vfs.VirtualFile; -import com.intellij.psi.PsiDocumentManager; -import com.intellij.psi.PsiElement; -import com.intellij.psi.PsiFile; -import com.intellij.psi.PsiReference; +import com.intellij.psi.*; import com.intellij.testFramework.PsiTestUtil; import com.jetbrains.python.psi.PyFile; import com.jetbrains.python.psi.PyFunction; @@ -34,6 +31,31 @@ public class PyMultiFileResolveTest extends CodeInsightTestCase { assertEquals("func", ((PyFunction) element).getName()); } + public void testFromPackageImport() throws Exception { + PsiElement element = doResolve(); + assertTrue(element instanceof PsiDirectory); + assertEquals("mypackage", ((PsiDirectory) element).getName()); + } + + public void testFromPackageImportFile() throws Exception { + PsiElement element = doResolve(); + assertTrue(element instanceof PsiFile); + assertEquals("myfile.py", ((PyFile) element).getName()); + } + + public void testFromQualifiedPackageImport() throws Exception { + PsiElement element = doResolve(); + assertTrue(element instanceof PsiDirectory); + assertEquals("mypackage", ((PsiDirectory) element).getName()); + } + + public void testFromQualifiedPackageImportFile() throws Exception { + PsiElement element = doResolve(); + assertTrue(element instanceof PsiFile); + assertEquals("myfile.py", ((PsiFile) element).getName()); + assertEquals("mypackage", ((PsiFile) element).getContainingDirectory().getName()); + } + private PsiElement doResolve() throws Exception { String testName = getTestName(true); String fileName = getTestName(false) + ".py";