one more case for stub-based resolve of qualified names

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -201,10 +201,13 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
       for (String component : qualifiedName.getComponents()) {
         PsiElement element = currentParent.getElementNamed(component);
         element = PyReferenceExpressionImpl.turnDirIntoInit(element);
+        if (element instanceof PyImportElement) {
+          element = ResolveImportUtil.resolveImportElement((PyImportElement) element);
+        }
         if (!(element instanceof NameDefiner)) {
           return null;
         }
-        currentParent = (NameDefiner) element;
+        currentParent = (NameDefiner)element;
       }
 
       if (!(currentParent instanceof PyClass)) {

python/src/com/jetbrains/python/psi/impl/PyFileImpl.java

@@ -188,25 +188,52 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
         if (child instanceof NamedStub && name.equals(((NamedStub)child).getName())) {
           return ((NamedStub) child).getPsi();
         }
-        else if (child instanceof PyFromImportStatementStub && ((PyFromImportStatementStub)child).isStarImport()) {
-          final PyFromImportStatement statement = ((PyFromImportStatementStub)child).getPsi();
-          PsiElement starImportSource = ResolveImportUtil.resolveFromImportStatementSource(statement);
-          if (starImportSource != null) {
-            starImportSource = PyReferenceExpressionImpl.turnDirIntoInit(starImportSource);
-            if (starImportSource instanceof PyFile) {
-              final PsiElement result = ((PyFile)starImportSource).getElementNamed(name);
-              if (result != null) {
-                return result;
+        else if (child instanceof PyFromImportStatementStub) {
+          if (((PyFromImportStatementStub)child).isStarImport()) {
+            final PyFromImportStatement statement = ((PyFromImportStatementStub)child).getPsi();
+            PsiElement starImportSource = ResolveImportUtil.resolveFromImportStatementSource(statement);
+            if (starImportSource != null) {
+              starImportSource = PyReferenceExpressionImpl.turnDirIntoInit(starImportSource);
+              if (starImportSource instanceof PyFile) {
+                final PsiElement result = ((PyFile)starImportSource).getElementNamed(name);
+                if (result != null) {
+                  return result;
+                }
+              }
+            }
+          }
+          else {
+            final List<StubElement> importElements = ((StubElement)child).getChildrenStubs();
+            for (StubElement importElement : importElements) {
+              final PsiElement psi = importElement.getPsi();
+              if (psi instanceof PyImportElement && name.equals(((PyImportElement)psi).getVisibleName())) {
+                return psi;
               }
             }
           }
         }
-        else if (child instanceof PyFromImportStatementStub || child instanceof PyImportStatementStub) {
+        else if (child instanceof PyImportStatementStub) {
           final List<StubElement> importElements = ((StubElement)child).getChildrenStubs();
-          for (StubElement importElement : importElements) {
-            final PsiElement psi = importElement.getPsi();
-            if (psi instanceof PyImportElement && name.equals(((PyImportElement)psi).getVisibleName())) {
-              return psi;
+          for (StubElement importElementStub : importElements) {
+            final PsiElement psi = importElementStub.getPsi();
+            if (psi instanceof PyImportElement) {
+              final PyImportElement importElement = (PyImportElement)psi;
+              final String asName = importElement.getAsName();
+              if (asName != null && asName.equals(name)) {
+                return psi;
+              }
+              final PyQualifiedName qName = importElement.getImportedQName();
+              if (qName != null && qName.getComponentCount() > 0) {
+                if (qName.getComponents().get(0).equals(name)) {
+                  if (qName.getComponentCount() == 1) {
+                    return psi;
+                  }
+                  return new PyImportedModule(this, PyQualifiedName.fromComponents(name));
+                }
+                if (name.equals(((PyImportElement)psi).getVisibleName())) {
+                  return psi;
+                }
+              }
             }
           }
         }

python/src/com/jetbrains/python/psi/impl/PyImportedModule.java

@@ -0,0 +1,80 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiElementVisitor;
+import com.intellij.psi.impl.light.LightElement;
+import com.jetbrains.python.PythonLanguage;
+import com.jetbrains.python.psi.NameDefiner;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.PyImportElement;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+
+/**
+ * @author yole
+ */
+public class PyImportedModule extends LightElement implements NameDefiner {
+  private final PyFile myContainingFile;
+  private final PyQualifiedName myImportedPrefix;
+
+  public PyImportedModule(PyFile containingFile, PyQualifiedName importedPrefix) {
+    super(containingFile.getManager(), PythonLanguage.getInstance());
+    myContainingFile = containingFile;
+    myImportedPrefix = importedPrefix;
+  }
+
+  @NotNull
+  public Iterable<PyElement> iterateNames() {
+    throw new UnsupportedOperationException();
+  }
+
+  public PsiElement getElementNamed(String the_name) {
+    PyQualifiedName prefix = myImportedPrefix.append(the_name);
+    final List<PyImportElement> imports = ((PyFileImpl)myContainingFile).getImportTargets();
+    for (PyImportElement anImport : imports) {
+      final PyQualifiedName qName = anImport.getImportedQName();
+      if (qName != null && matchesPrefix(qName, prefix)) {
+        if (qName.getComponentCount() == prefix.getComponentCount()) {
+          return anImport;
+        }
+        return new PyImportedModule(myContainingFile, prefix);
+      }
+    }
+    return null;
+  }
+
+  private static boolean matchesPrefix(PyQualifiedName qName, PyQualifiedName prefix) {
+    if (qName.getComponentCount() < prefix.getComponentCount()) {
+      return false;
+    }
+    for (int i = 0; i < prefix.getComponentCount(); i++) {
+      if (!qName.getComponents().get(i).equals(prefix.getComponents().get(i))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  public boolean mustResolveOutside() {
+    return true;
+  }
+
+  public String getText() {
+    return "import " + myImportedPrefix;
+  }
+
+  public void accept(@NotNull PsiElementVisitor visitor) {
+    visitor.visitElement(this);
+  }
+
+  public PsiElement copy() {
+    return new PyImportedModule(myContainingFile, myImportedPrefix);
+  }
+
+  @Override
+  public String toString() {
+    return "PyImportedModule:" + myImportedPrefix;
+  }
+}

python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java

@@ -7,6 +7,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -26,6 +27,19 @@ public class PyQualifiedName {
     }
   }
 
+  public static PyQualifiedName fromComponents(String... components) {
+    PyQualifiedName result = new PyQualifiedName(components.length);
+    Collections.addAll(result.myComponents, components);
+    return result;
+  }
+
+  public PyQualifiedName append(String name) {
+    PyQualifiedName result = new PyQualifiedName(myComponents.size()+1);
+    result.myComponents.addAll(myComponents);
+    result.myComponents.add(name);
+    return result;
+  }
+
   public List<String> getComponents() {
     return myComponents;
   }

python/testData/resolve/multiFile/resolveQualifiedSuperClassInPackage/ResolveQualifiedSuperClassInPackage.py

@@ -0,0 +1,6 @@
+import bar
+
+class Foo(bar.Super): pass
+
+Foo().copy()
+#      <ref>

python/testData/resolve/multiFile/resolveQualifiedSuperClassInPackage/bar/__init__.py

@@ -0,0 +1,5 @@
+import foo.baz
+
+class Super(foo.baz.SuperDuper):
+  pass
+

python/testData/resolve/multiFile/resolveQualifiedSuperClassInPackage/foo/baz.py

@@ -0,0 +1,2 @@
+class SuperDuper(object):
+    def copy(self): print "duper"

python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java

@@ -154,6 +154,12 @@ public class PyMultiFileResolveTest extends PyLightFixtureTestCase {
     assertEquals("copy", ((PyFunction) element).getName());
   }
 
+  public void testResolveQualifiedSuperClassInPackage() throws Exception {
+    PsiElement element = doResolve();
+    assertTrue(element instanceof PyFunction);
+    assertEquals("copy", ((PyFunction) element).getName());
+  }
+
   private PsiFile prepareFile() throws Exception {
     String testName = getTestName(true);
     String fileName = getTestName(false) + ".py";