resolving imports cont'd

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);
     }
 

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;
+    }
+  }
 }

python/testData/resolve/multiFile/fromPackageImport/FromPackageImport.py

@@ -0,0 +1,2 @@
+from mypackage import p
+#     <ref>

python/testData/resolve/multiFile/fromPackageImportFile/FromPackageImportFile.py

@@ -0,0 +1,2 @@
+from mypackage import myfile
+#                      <ref>

python/testData/resolve/multiFile/fromQualifiedPackageImport/FromQualifiedPackageImport.py

@@ -0,0 +1,2 @@
+from mypackage.myfile import myclass
+#      <ref>

python/testData/resolve/multiFile/fromQualifiedPackageImport/mypackage/myfile.py

@@ -0,0 +1,2 @@
+class myclass:
+  pass

python/testData/resolve/multiFile/fromQualifiedPackageImportFile/FromQualifiedPackageImportFile.py

@@ -0,0 +1,2 @@
+from mypackage.myfile import myclass
+#               <ref>

python/testData/resolve/multiFile/fromQualifiedPackageImportFile/mypackage/myfile.py

@@ -0,0 +1,2 @@
+class myclass:
+  pass

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";