separate implementation class for references inside import statements

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

@@ -0,0 +1,44 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.PyUtil;
+import com.jetbrains.python.psi.RatedResolveResult;
+import com.jetbrains.python.psi.resolve.ResolveImportUtil;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+
+/**
+ * @author yole
+ */
+public class PyImportReferenceImpl extends PyReferenceImpl {
+  public PyImportReferenceImpl(PyReferenceExpressionImpl element) {
+    super(element);
+  }
+
+  @NotNull
+  @Override
+  protected List<RatedResolveResult> resolveInner() {
+    ResultList ret = new ResultList();
+
+    final String referencedName = myElement.getReferencedName();
+    if (referencedName == null) return ret;
+
+    PsiElement target = ResolveImportUtil.resolveImportReference(myElement);
+
+    target = PyUtil.turnDirIntoInit(target);
+    if (target == null) {
+      ret.clear();
+      return ret; // it was a dir without __init__.py, worthless
+    }
+    ret.poke(target, RatedResolveResult.RATE_HIGH);
+    return ret;
+  }
+
+  @NotNull
+  @Override
+  public Object[] getVariants() {
+    // complete to possible modules
+    return ResolveImportUtil.suggestImportVariants(myElement);
+  }
+}

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

@@ -1,7 +1,6 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.psi.*;
@@ -28,16 +27,19 @@ import java.util.List;
  */
 public class PyReferenceExpressionImpl extends PyElementImpl implements PyReferenceExpression {
   private static final Logger LOG = Logger.getInstance("#com.jetbrains.python.psi.impl.PyReferenceExpressionImpl");
-  private final PyReferenceImpl myReference;
 
   public PyReferenceExpressionImpl(ASTNode astNode) {
     super(astNode);
-    myReference = new PyReferenceImpl(this);
   }
 
   @NotNull
   public PsiPolyVariantReference getReference() {
-    return myReference;
+    // Handle import reference
+    if (PsiTreeUtil.getParentOfType(this, PyImportElement.class, PyFromImportStatement.class) != null) {
+      return new PyImportReferenceImpl(this);
+    }
+
+    return new PyReferenceImpl(this);
   }
 
   @Override
@@ -75,7 +77,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     PyElement ret = null;
     SEARCH:
     while (ret == null) {
-      ResolveResult[] targets = ((PsiPolyVariantReference) seeker.getReference()).multiResolve(false);
+      ResolveResult[] targets = seeker.getReference().multiResolve(false);
       for (ResolveResult target : targets) {
         PsiElement elt = target.getElement();
         if (elt instanceof PyTargetExpression) {
@@ -151,7 +153,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
       return pyType;
     }
 
-    ResolveResult[] targets = myReference.multiResolve(false);
+    ResolveResult[] targets = getReference().multiResolve(false);
     if (targets.length == 0) return null;
     PsiElement target = targets[0].getElement();
     if (target == this) {

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

@@ -30,7 +30,7 @@ import java.util.*;
  * @author yole
  */
 public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference {
-  private final PyReferenceExpressionImpl myElement;
+  protected final PyReferenceExpressionImpl myElement;
 
   public PyReferenceImpl(PyReferenceExpressionImpl element) {
     myElement = element;
@@ -142,7 +142,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     return ret.toArray(new ResolveResult[ret.size()]);
   }
 
-  private static class ResultList extends ArrayList<RatedResolveResult> {
+  protected static class ResultList extends ArrayList<RatedResolveResult> {
     // Allows to add non-null elements and discard nulls in a hassle-free way.
     public boolean poke(final PsiElement what, final int rate) {
       if (what == null) return false;
@@ -167,7 +167,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
    * @see #resolve()
    */
   @NotNull
-  private List<RatedResolveResult> resolveInner() {
+  protected List<RatedResolveResult> resolveInner() {
     ResultList ret = new ResultList();
 
     final String referencedName = myElement.getReferencedName();
@@ -175,15 +175,6 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
 
     // Handle import reference
     if (PsiTreeUtil.getParentOfType(myElement, PyImportElement.class, PyFromImportStatement.class) != null) {
-      PsiElement target = ResolveImportUtil.resolveImportReference(myElement);
-
-      target = PyUtil.turnDirIntoInit(target);
-      if (target == null) {
-        ret.clear();
-        return ret; // it was a dir without __init__.py, worthless
-      }
-      ret.poke(target, RatedResolveResult.RATE_HIGH);
-      return ret;
     }
 
     final PyExpression qualifier = myElement.getQualifier();
@@ -349,12 +340,6 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
 
   @NotNull
   public Object[] getVariants() {
-    // imports are another special case
-    if (PsiTreeUtil.getParentOfType(myElement, PyImportElement.class, PyFromImportStatement.class) != null) {
-      // complete to possible modules
-      return ResolveImportUtil.suggestImportVariants(myElement);
-    }
-
     // qualifier limits the namespace
     final PyExpression qualifier = myElement.getQualifier();
     if (qualifier != null) {