Use multi-resolve for multiple values conditionally defined in an imported module (PY-18402)

Added PyFile.multiResolveName() that is a multi-resolve version of
getElementNamed(). The same method is added to PyImportedNameDefiner,
but there is no need in extracting a super inferface at the moment.
Finally, there is the new PyImportElement.multiResolve() in addition to
resolve().

Changed some single-resolve usages to multi-resolve acorss the codebase.
There are still many less important single-resolve usages that should
be converted to multi-resolve eventually. Single-resolve variants are
deprecated now.

python/psi-api/src/com/jetbrains/python/psi/PyFile.java

@@ -18,6 +18,7 @@ package com.jetbrains.python.psi;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -65,8 +66,15 @@ public interface PyFile extends PyElement, PsiFile, PyDocStringOwner, ScopeOwner
   Iterable<PyElement> iterateNames();
 
   /**
-   * Return the resolved exported PSI element with the top priority among resolve targets.
+   * Return the resolved exported elements.
    */
+  @NotNull
+  List<RatedResolveResult> multiResolveName(@NotNull String name);
+
+  /**
+   * @deprecated Use {@link #multiResolveName(String)} instead.
+   */
+  @Deprecated
   @Nullable
   PsiElement getElementNamed(String name);
 

python/psi-api/src/com/jetbrains/python/psi/PyImportElement.java

@@ -18,9 +18,13 @@ package com.jetbrains.python.psi;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.StubBasedPsiElement;
 import com.intellij.psi.util.QualifiedName;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.stubs.PyImportElementStub;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
+
 /**
  * @author yole
  */
@@ -49,10 +53,15 @@ public interface PyImportElement extends PyElement, PyImportedNameDefiner, StubB
   PsiElement getElementNamed(String name, boolean resolveImportElement);
 
   /**
-   * Resolves the import element to the element being imported.
-   *
-   * @return the resolve result or null if the resolution failed.
+   * @deprecated Use {@link #multiResolve()} instead.
    */
+  @Deprecated
   @Nullable
   PsiElement resolve();
+
+  /**
+   * Resolves the import element to the elements being imported.
+   */
+  @NotNull
+  List<RatedResolveResult> multiResolve();
 }

python/psi-api/src/com/jetbrains/python/psi/PyImportedNameDefiner.java

@@ -15,9 +15,10 @@
  */
 package com.jetbrains.python.psi;
 
-import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
 
 /**
  * Name definer that defines names imported somehow from other modules.
@@ -26,7 +27,9 @@ import org.jetbrains.annotations.Nullable;
  */
 public interface PyImportedNameDefiner extends PyElement {
   /**
-   * Iterate over resolved PSI elements available via this imported name definer.
+   * Iterate over possibly resolved PSI elements available via this imported name definer.
+   *
+   * TODO: Make the semantics of the returned elements clearer.
    */
   @NotNull
   Iterable<PyElement> iterateNames();
@@ -34,6 +37,6 @@ public interface PyImportedNameDefiner extends PyElement {
   /**
    * Return the resolved PSI element available via this imported name definer.
    */
-  @Nullable
-  PsiElement getElementNamed(String name);
+  @NotNull
+  List<RatedResolveResult> multiResolveName(@NotNull String name);
 }

python/src/com/jetbrains/python/codeInsight/dataflow/scope/impl/ScopeImpl.java

@@ -126,7 +126,7 @@ public class ScopeImpl implements Scope {
       return true;
     }
     for (PyImportedNameDefiner definer : getImportedNameDefiners()) {
-      if (definer.getElementNamed(name) != null) {
+      if (!definer.multiResolveName(name).isEmpty()) {
         return true;
       }
     }

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

@@ -16,6 +16,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.intellij.extapi.psi.PsiFileBase;
 import com.intellij.icons.AllIcons;
 import com.intellij.lang.ASTNode;
@@ -39,7 +40,6 @@ import com.intellij.reference.SoftReference;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.Processor;
 import com.intellij.util.containers.ContainerUtil;
-import com.intellij.util.containers.MultiMap;
 import com.intellij.util.indexing.IndexingDataKeys;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
@@ -76,8 +76,8 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
     private final List<String> myNameDefinerNegativeCache = new ArrayList<String>();
     private long myNameDefinerOOCBModCount = -1;
     private final long myModificationStamp;
-    private final MultiMap<String, PsiNamedElement> myNamedElements = new MultiMap<String, PsiNamedElement>();
-    private final List<PyImportedNameDefiner> myImportedNameDefiners = new ArrayList<PyImportedNameDefiner>();
+    private final Map<String, List<PsiNamedElement>> myNamedElements = Maps.newHashMap();
+    private final List<PyImportedNameDefiner> myImportedNameDefiners = Lists.newArrayList();
 
     private ExportedNameCache(long modificationStamp) {
       myModificationStamp = modificationStamp;
@@ -87,7 +87,12 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
         public boolean process(PsiElement element) {
           if (element instanceof PsiNamedElement && !(element instanceof PyKeywordArgument)) {
             final PsiNamedElement namedElement = (PsiNamedElement)element;
-            myNamedElements.putValue(namedElement.getName(), namedElement);
+            final String name = namedElement.getName();
+            if (!myNamedElements.containsKey(name)) {
+              myNamedElements.put(name, Lists.<PsiNamedElement>newArrayList());
+            }
+            final List<PsiNamedElement> elements = myNamedElements.get(name);
+            elements.add(namedElement);
           }
           if (element instanceof PyImportedNameDefiner) {
             myImportedNameDefiners.add((PyImportedNameDefiner)element);
@@ -109,6 +114,10 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
           return true;
         }
       });
+      for (List<PsiNamedElement> elements : myNamedElements.values()) {
+        Collections.reverse(elements);
+      }
+      Collections.reverse(myImportedNameDefiners);
     }
 
     private boolean processDeclarations(@NotNull List<PsiElement> elements, @NotNull Processor<PsiElement> processor) {
@@ -381,7 +390,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
   @Override
   @Nullable
   public PsiElement findExportedName(final String name) {
-    final List<RatedResolveResult> results = multiResolveExportedElements(name);
+    final List<RatedResolveResult> results = multiResolveName(name);
     final List<PsiElement> elements = Lists.newArrayList();
     for (RatedResolveResult result : results) {
       final PsiElement element = result.getElement();
@@ -404,14 +413,22 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
   }
 
   @NotNull
-  private List<RatedResolveResult> multiResolveExportedElements(@NotNull final String name) {
+  @Override
+  public List<RatedResolveResult> multiResolveName(@NotNull final String name) {
     final List<RatedResolveResult> results = RecursionManager.doPreventingRecursion(this, false, new Computable<List<RatedResolveResult>>() {
       @Override
       public List<RatedResolveResult> compute() {
         return getExportedNameCache().multiResolve(name);
       }
     });
-    return results != null ? results : Collections.<RatedResolveResult>emptyList();
+    if (results != null && !results.isEmpty()) {
+      return results;
+    }
+    final List<String> allNames = getDunderAll();
+    if (allNames != null && allNames.contains(name)) {
+      return ResolveResultList.to(findExportedName(PyNames.ALL));
+    }
+    return Collections.emptyList();
   }
 
   private ExportedNameCache getExportedNameCache() {
@@ -431,7 +448,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
 
   @Nullable
   public PsiElement getElementNamed(final String name) {
-    final List<RatedResolveResult> results = multiResolveExportedElements(name);
+    final List<RatedResolveResult> results = multiResolveName(name);
     final List<PsiElement> elements = PyUtil.filterTopPriorityResults(results.toArray(new ResolveResult[results.size()]));
     final PsiElement element = elements.isEmpty() ? null : elements.get(elements.size() - 1);
     if (element != null) {
@@ -440,10 +457,6 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
       }
       return element;
     }
-    final List<String> allNames = getDunderAll();
-    if (allNames != null && allNames.contains(name)) {
-      return findExportedName(PyNames.ALL);
-    }
     return null;
   }
 

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

@@ -29,6 +29,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.resolve.ResolveImportUtil;
 import com.jetbrains.python.psi.stubs.PyFromImportStatementStub;
 import org.jetbrains.annotations.NotNull;
@@ -263,11 +264,17 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     return resolved != null ? ImmutableList.<PyElement>of(resolved) : Collections.<PyElement>emptyList();
   }
 
-  @Nullable
+  @NotNull
   @Override
-  public PsiElement getElementNamed(String name) {
+  public List<RatedResolveResult> multiResolveName(@NotNull String name) {
     final QualifiedName importSourceQName = getImportSourceQName();
-    return importSourceQName != null && importSourceQName.endsWith(name) ? resolveImplicitSubModule() : null;
+    if (importSourceQName != null && importSourceQName.endsWith(name)) {
+      final PsiElement element = resolveImplicitSubModule();
+      if (element != null) {
+        return ResolveResultList.to(element);
+      }
+    }
+    return Collections.emptyList();
   }
 
   /**

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

@@ -28,6 +28,7 @@ import com.intellij.util.containers.EmptyIterable;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.resolve.ResolveImportUtil;
 import com.jetbrains.python.psi.stubs.PyImportElementStub;
 import org.jetbrains.annotations.NotNull;
@@ -35,6 +36,7 @@ import org.jetbrains.annotations.Nullable;
 
 import javax.swing.*;
 import java.util.Collections;
+import java.util.List;
 
 /**
  * The "import foo" or "import foo as bar" parts.
@@ -193,37 +195,57 @@ public class PyImportElementImpl extends PyBaseElementImpl<PyImportElementStub>
     return Collections.singleton(ret);
   }
 
-  public PsiElement getElementNamed(final String name) {
-    return getElementNamed(name, true);
+  @NotNull
+  public List<RatedResolveResult> multiResolveName(@NotNull final String name) {
+    return getElementsNamed(name, true);
   }
 
+  @Nullable
   @Override
   public PsiElement getElementNamed(String name, boolean resolveImportElement) {
+    final List<RatedResolveResult> results = getElementsNamed(name, resolveImportElement);
+    return results.isEmpty() ? null : RatedResolveResult.sorted(results).get(0).getElement();
+  }
+
+  @NotNull
+  private List<RatedResolveResult> getElementsNamed(@NotNull String name, boolean resolveImportElement) {
     String asName = getAsName();
     if (asName != null) {
-      if (!Comparing.equal(name, asName)) return null;
-      return resolveImportElement ? resolve() : this;
+      if (!Comparing.equal(name, asName)) {
+        return Collections.emptyList();
+      }
+      if (resolveImportElement) {
+        return multiResolve();
+      }
+      return ResolveResultList.to(this);
     }
     else {
       final QualifiedName qName = getImportedQName();
       if (qName == null || qName.getComponentCount() == 0 || !qName.getComponents().get(0).equals(name)) {
-        return null;
+        return Collections.emptyList();
       }
       if (qName.getComponentCount() == 1) {
         if (resolveImportElement) {
-          return resolve();
+          return multiResolve();
         }
-        return this;
+        return ResolveResultList.to(this);
       }
-      return createImportedModule(name);
+      return ResolveResultList.to(createImportedModule(name));
     }
   }
 
   @Nullable
   @Override
   public PsiElement resolve() {
-    QualifiedName qName = getImportedQName();
-    return qName == null ? null : ResolveImportUtil.resolveImportElement(this, qName);
+    final List<RatedResolveResult> results = multiResolve();
+    return results.isEmpty() ? null : RatedResolveResult.sorted(results).get(0).getElement();
+  }
+
+  @NotNull
+  @Override
+  public List<RatedResolveResult> multiResolve() {
+    final QualifiedName qName = getImportedQName();
+    return qName == null ? Collections.<RatedResolveResult>emptyList() : ResolveImportUtil.multiResolveImportElement(this, qName);
   }
 
   @Override

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

@@ -26,6 +26,7 @@ import com.intellij.util.ArrayFactory;
 import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.resolve.ResolveImportUtil;
 import com.jetbrains.python.psi.stubs.PyImportStatementStub;
 import org.jetbrains.annotations.NotNull;
@@ -110,18 +111,18 @@ public class PyImportStatementImpl extends PyBaseElementImpl<PyImportStatementSt
     return resolved != null ? ImmutableList.<PyElement>of(resolved) : Collections.<PyElement>emptyList();
   }
 
-  @Nullable
+  @NotNull
   @Override
-  public PsiElement getElementNamed(String name) {
+  public List<RatedResolveResult> multiResolveName(@NotNull String name) {
     final PyImportElement[] elements = getImportElements();
     if (elements.length == 1) {
       final PyImportElement element = elements[0];
       final QualifiedName importedQName = element.getImportedQName();
       if (importedQName != null && importedQName.getComponentCount() > 1 && name.equals(importedQName.getLastComponent())) {
-        return resolveImplicitSubModule();
+        return ResolveResultList.to(resolveImplicitSubModule());
       }
     }
-    return null;
+    return Collections.emptyList();
   }
 
   /**

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

@@ -17,6 +17,7 @@ package com.jetbrains.python.psi.impl;
 
 import com.google.common.base.Predicate;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
 import com.intellij.lang.ASTNode;
 import com.intellij.navigation.ItemPresentation;
 import com.intellij.psi.PsiElement;
@@ -78,10 +79,10 @@ public class PyStarImportElementImpl extends PyBaseElementImpl<PyStarImportEleme
     });
   }
 
-  @Nullable
-  public PsiElement getElementNamed(final String name) {
+  @NotNull
+  public List<RatedResolveResult> multiResolveName(@NotNull final String name) {
     if (PyUtil.isClassPrivateName(name)) {
-      return null;
+      return Collections.emptyList();
     }
     final PsiElement parent = getParentByStub();
     if (parent instanceof PyFromImportStatement) {
@@ -94,14 +95,17 @@ public class PyStarImportElementImpl extends PyBaseElementImpl<PyStarImportEleme
           final PyModuleType moduleType = new PyModuleType(sourceFile);
           final List<? extends RatedResolveResult> results = moduleType.resolveMember(name, null, AccessDirection.READ,
                                                                                       PyResolveContext.defaultContext());
-          final PsiElement result = results != null && !results.isEmpty() ? results.get(0).getElement() : null;
-          if (result != null && PyUtil.isStarImportableFrom(name, sourceFile) ) {
-            return result;
+          if (results != null && !results.isEmpty() && PyUtil.isStarImportableFrom(name, sourceFile)) {
+            final List<RatedResolveResult> res = Lists.newArrayList();
+            for (RatedResolveResult result : results) {
+              res.add(result);
+            }
+            return res;
           }
         }
       }
     }
-    return null;
+    return Collections.emptyList();
   }
 
   @Override

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

@@ -165,18 +165,21 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
                                                        @NotNull TypeEvalContext context) {
     final ResolveResultList ret = new ResolveResultList();
     for (Instruction instruction : instructions) {
-      PsiElement definition = instruction.getElement();
-      PyImportedNameDefiner definer = null;
+      final PsiElement definition = instruction.getElement();
       // TODO: This check may slow down resolving, but it is the current solution to the comprehension scopes problem
       if (isInnerComprehension(element, definition)) continue;
       if (definition instanceof PyImportedNameDefiner && !(definition instanceof PsiNamedElement)) {
-        definer = (PyImportedNameDefiner)definition;
-        definition = definer.getElementNamed(name);
-      }
-      if (definer != null) {
-        ret.add(new ImportedResolveResult(definition, getRate(definition, context), definer));
-        // TODO this kind of resolve contract is quite stupid
-        if (definition != null) {
+        final PyImportedNameDefiner definer = (PyImportedNameDefiner)definition;
+        final List<RatedResolveResult> resolvedResults = definer.multiResolveName(name);
+        for (RatedResolveResult result : resolvedResults) {
+          final PsiElement resolved = result.getElement();
+          ret.add(new ImportedResolveResult(resolved, getRate(resolved, context), definer));
+        }
+        if (resolvedResults.isEmpty()) {
+          ret.add(new ImportedResolveResult(null, RatedResolveResult.RATE_NORMAL, definer));
+        }
+        else {
+          // TODO this kind of resolve contract is quite stupid
           ret.poke(definer, RatedResolveResult.RATE_LOW);
         }
       }

python/src/com/jetbrains/python/psi/resolve/PyResolveProcessor.java

@@ -24,10 +24,13 @@ import com.intellij.psi.scope.PsiScopeProcessor;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.ResolveResultList;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -57,9 +60,16 @@ public class PyResolveProcessor implements PsiScopeProcessor {
     }
     final PyImportedNameDefiner importedNameDefiner = PyUtil.as(element, PyImportedNameDefiner.class);
     if (importedNameDefiner != null) {
-      final PsiElement resolved = resolveInImportedNameDefiner(importedNameDefiner);
-      if (resolved != null) {
-        return tryAddResult(resolved, importedNameDefiner);
+      final List<RatedResolveResult> results = resolveInImportedNameDefiner(importedNameDefiner);
+      if (!results.isEmpty()) {
+        boolean cont = true;
+        for (RatedResolveResult result : results) {
+          final PsiElement resolved = result.getElement();
+          if (resolved != null) {
+            cont = tryAddResult(resolved, importedNameDefiner) && cont;
+          }
+        }
+        return cont;
       }
       final PyImportElement importElement = PyUtil.as(element, PyImportElement.class);
       if (importElement != null) {
@@ -97,18 +107,18 @@ public class PyResolveProcessor implements PsiScopeProcessor {
     return myOwner;
   }
 
-  @Nullable
-  private PsiElement resolveInImportedNameDefiner(@NotNull PyImportedNameDefiner definer) {
+  @NotNull
+  private List<RatedResolveResult> resolveInImportedNameDefiner(@NotNull PyImportedNameDefiner definer) {
     if (myLocalResolve) {
       final PyImportElement importElement = PyUtil.as(definer, PyImportElement.class);
       if (importElement != null) {
-        return importElement.getElementNamed(myName, false);
+        return ResolveResultList.to(importElement.getElementNamed(myName, false));
       }
       else if (definer instanceof PyStarImportElement) {
-        return null;
+        return Collections.emptyList();
       }
     }
-    return definer.getElementNamed(myName);
+    return definer.multiResolveName(myName);
   }
 
   private boolean tryAddResult(@Nullable PsiElement element, @Nullable PyImportedNameDefiner definer) {

python/src/com/jetbrains/python/psi/resolve/QualifiedNameResolverImpl.java

@@ -72,6 +72,9 @@ public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameReso
     myQualifiedName = QualifiedName.fromDottedString(qNameString);
   }
 
+  /**
+   * @param qName the empty name means that all found roots will be traversed.
+   */
   public QualifiedNameResolverImpl(@NotNull QualifiedName qName) {
     myQualifiedName = qName;
   }
@@ -253,8 +256,10 @@ public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameReso
       addRelativeImportResultsFromSkeletons(footholdFile);
     }
 
-    mySourceResults.addAll(myLibResults);
-    myLibResults.clear();
+    if (mySourceResults.isEmpty() || myQualifiedName.getComponentCount() == 0) {
+      mySourceResults.addAll(myLibResults);
+      myLibResults.clear();
+    }
 
     if (!myWithoutForeign) {
       for (PyImportResolver resolver : Extensions.getExtensions(PyImportResolver.EP_NAME)) {
@@ -263,8 +268,10 @@ public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameReso
           myForeignResults.add(foreign);
         }
       }
-      mySourceResults.addAll(myForeignResults);
-      myForeignResults.clear();
+      if (mySourceResults.isEmpty() || myQualifiedName.getComponentCount() == 0) {
+        mySourceResults.addAll(myForeignResults);
+        myForeignResults.clear();
+      }
     }
 
     final ArrayList<PsiElement> results = Lists.newArrayList(mySourceResults);

python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java

@@ -15,6 +15,7 @@
  */
 package com.jetbrains.python.psi.resolve;
 
+import com.google.common.collect.Lists;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.fileTypes.ExtensionFileNameMatcher;
 import com.intellij.openapi.fileTypes.FileNameMatcher;
@@ -104,20 +105,28 @@ public class ResolveImportUtil {
     return null;
   }
 
+  /**
+   * @deprecated use {@link #multiResolveImportElement(PyImportElement, QualifiedName)} instead.
+   */
+  @Deprecated
   @Nullable
   public static PsiElement resolveImportElement(PyImportElement importElement, @NotNull final QualifiedName qName) {
-    List<RatedResolveResult> targets;
+    final List<RatedResolveResult> resultList = RatedResolveResult.sorted(multiResolveImportElement(importElement, qName));
+    return resultList.size() > 0 ? resultList.get(0).getElement() : null;
+  }
+
+  @NotNull
+  public static List<RatedResolveResult> multiResolveImportElement(PyImportElement importElement, @NotNull final QualifiedName qName) {
     final PyStatement importStatement = importElement.getContainingImportStatement();
     if (importStatement instanceof PyFromImportStatement) {
-      targets = resolveNameInFromImport((PyFromImportStatement)importStatement, qName);
+      return resolveNameInFromImport((PyFromImportStatement)importStatement, qName);
     }
-    else { // "import foo"
-      targets = resolveNameInImportStatement(importElement, qName);
+    else {
+      return resolveNameInImportStatement(importElement, qName);
     }
-    final List<RatedResolveResult> resultList = RatedResolveResult.sorted(targets);
-    return resultList.size() > 0 ? resultList.get(0).getElement() : null;
   }
 
+  @NotNull
   public static List<RatedResolveResult> resolveNameInImportStatement(PyImportElement importElement, @NotNull QualifiedName qName) {
     final PsiFile file = importElement.getContainingFile().getOriginalFile();
     boolean absoluteImportEnabled = isAbsoluteImportEnabledFor(importElement);
@@ -125,6 +134,7 @@ public class ResolveImportUtil {
     return rateResults(modules);
   }
 
+  @NotNull
   public static List<RatedResolveResult> resolveNameInFromImport(PyFromImportStatement importStatement, @NotNull QualifiedName qName) {
     PsiFile file = importStatement.getContainingFile().getOriginalFile();
     String name = qName.getComponents().get(0);
@@ -138,12 +148,17 @@ public class ResolveImportUtil {
       if (candidate instanceof PsiDirectory) {
         candidate = PyUtil.getPackageElement((PsiDirectory)candidate, importStatement);
       }
-      PsiElement result = resolveChild(candidate, name, file, false, true);
-      if (result != null) {
-        if (!result.isValid()) {
-          throw new PsiInvalidElementAccessException(result, "Got an invalid candidate from resolveChild(): " + result.getClass());
+      List<RatedResolveResult> results = resolveChildren(candidate, name, file, false, true);
+      if (!results.isEmpty()) {
+        for (RatedResolveResult result : results) {
+          final PsiElement element = result.getElement();
+          if (element != null) {
+            if (!element.isValid()) {
+              throw new PsiInvalidElementAccessException(element, "Got an invalid candidate from resolveChild(): " + element.getClass());
+            }
+            resultList.add(element);
+          }
         }
-        resultList.add(result);
       }
     }
     if (!resultList.isEmpty()) {
@@ -256,6 +271,17 @@ public class ResolveImportUtil {
     return cache;
   }
 
+  /**
+   * @deprecated Use {@link #resolveChildren(PsiElement, String, PsiFile, boolean, boolean)} instead.
+   */
+  @Deprecated
+  @Nullable
+  public static PsiElement resolveChild(@Nullable final PsiElement parent, @NotNull final String referencedName,
+                                        @Nullable final PsiFile containingFile, boolean fileOnly, boolean checkForPackage) {
+    final List<RatedResolveResult> results = resolveChildren(parent, referencedName, containingFile, fileOnly, checkForPackage);
+    return results.isEmpty() ? null : RatedResolveResult.sorted(results).get(0).getElement();
+  }
+
   /**
    * Tries to find referencedName under the parent element.
    *
@@ -266,11 +292,11 @@ public class ResolveImportUtil {
    * @param checkForPackage if true, directories are returned only if they contain __init__.py
    * @return the element the referencedName resolves to, or null.
    */
-  @Nullable
-  public static PsiElement resolveChild(@Nullable final PsiElement parent, @NotNull final String referencedName,
-                                        @Nullable final PsiFile containingFile, boolean fileOnly, boolean checkForPackage) {
+  @NotNull
+  public static List<RatedResolveResult> resolveChildren(@Nullable PsiElement parent, @NotNull String referencedName,
+                                                         @Nullable PsiFile containingFile, boolean fileOnly, boolean checkForPackage) {
     if (parent == null) {
-      return null;
+      return Collections.emptyList();
     }
     else if (parent instanceof PyFile) {
       return resolveInPackageModule((PyFile)parent, referencedName, containingFile, fileOnly, checkForPackage);
@@ -281,24 +307,36 @@ public class ResolveImportUtil {
     else {
       return resolveMemberFromReferenceTypeProviders(parent, referencedName);
     }
+
   }
 
-  @Nullable
-  private static PsiElement resolveInPackageModule(@NotNull PyFile parent, @NotNull String referencedName,
-                                                   @Nullable PsiFile containingFile, boolean fileOnly, boolean checkForPackage) {
-    final PsiElement moduleMember = resolveModuleMember(parent, referencedName);
-    final PsiElement resolved = !fileOnly || PyUtil.instanceOf(moduleMember, PsiFile.class, PsiDirectory.class) ?
-                                moduleMember : null;
-    if (resolved != null && !preferResolveInDirectoryOverModule(resolved)) {
-      return resolved;
+  @NotNull
+  private static List<RatedResolveResult> resolveInPackageModule(@NotNull PyFile parent, @NotNull String referencedName,
+                                                                 @Nullable PsiFile containingFile, boolean fileOnly,
+                                                                 boolean checkForPackage) {
+    final List<RatedResolveResult> moduleMembers = resolveModuleMember(parent, referencedName);
+    final List<RatedResolveResult> resolvedInModule = Lists.newArrayList();
+    final List<RatedResolveResult> results = Lists.newArrayList();
+    for (RatedResolveResult member : moduleMembers) {
+      final PsiElement moduleMember = member.getElement();
+      if (!fileOnly || PyUtil.instanceOf(moduleMember, PsiFile.class, PsiDirectory.class)) {
+        results.add(member);
+        if (moduleMember != null && !preferResolveInDirectoryOverModule(moduleMember)) {
+          resolvedInModule.add(member);
+        }
+      }
+    }
+    if (!resolvedInModule.isEmpty()) {
+      return resolvedInModule;
     }
 
-    final PsiElement resolvedInDirectory = resolveInPackageDirectory(parent, referencedName, containingFile, fileOnly, checkForPackage);
-    if (resolvedInDirectory != null) {
+    final List<RatedResolveResult> resolvedInDirectory = resolveInPackageDirectory(parent, referencedName, containingFile, fileOnly,
+                                                                                   checkForPackage);
+    if (!resolvedInDirectory.isEmpty()) {
       return resolvedInDirectory;
     }
 
-    return resolved;
+    return results;
   }
 
   private static boolean preferResolveInDirectoryOverModule(@NotNull PsiElement resolved) {
@@ -307,30 +345,34 @@ public class ResolveImportUtil {
            isDunderAll(resolved);
   }
 
-  @Nullable
-  private static PsiElement resolveModuleMember(@NotNull PyFile file, @NotNull String referencedName) {
+  @NotNull
+  private static List<RatedResolveResult> resolveModuleMember(@NotNull PyFile file, @NotNull String referencedName) {
     final PyModuleType moduleType = new PyModuleType(file);
     final PyResolveContext resolveContext = PyResolveContext.defaultContext();
     final List<? extends RatedResolveResult> results = moduleType.resolveMember(referencedName, null, AccessDirection.READ,
                                                                                 resolveContext);
-    return results != null && !results.isEmpty() ? results.get(0).getElement() : null;
+    if (results == null) {
+      return Collections.emptyList();
+    }
+    return Lists.newArrayList(results);
   }
 
-  @Nullable
-  private static PsiElement resolveInPackageDirectory(@Nullable PsiElement parent, @NotNull String referencedName,
-                                                      @Nullable PsiFile containingFile, boolean fileOnly,
-                                                      boolean checkForPackage) {
+  @NotNull
+  private static List<RatedResolveResult> resolveInPackageDirectory(@Nullable PsiElement parent, @NotNull String referencedName,
+                                                                    @Nullable PsiFile containingFile, boolean fileOnly,
+                                                                    boolean checkForPackage) {
     final PsiElement parentDir = PyUtil.turnInitIntoDir(parent);
     if (parentDir instanceof PsiDirectory) {
-      final PsiElement resolved = resolveInDirectory(referencedName, containingFile, (PsiDirectory)parentDir, fileOnly, checkForPackage);
-      if (resolved != null) {
+      final List<RatedResolveResult> resolved = resolveInDirectory(referencedName, containingFile, (PsiDirectory)parentDir, fileOnly,
+                                                                   checkForPackage);
+      if (!resolved.isEmpty()) {
         return resolved;
       }
       if (parent instanceof PsiFile) {
-        return resolveForeignImports((PsiFile)parent, referencedName);
+        return ResolveResultList.to(resolveForeignImports((PsiFile)parent, referencedName));
       }
     }
-    return null;
+    return Collections.emptyList();
   }
 
   @Nullable
@@ -338,45 +380,48 @@ public class ResolveImportUtil {
     return new QualifiedNameResolverImpl(referencedName).fromElement(foothold).withoutRoots().firstResult();
   }
 
-  @Nullable
-  private static PsiElement resolveMemberFromReferenceTypeProviders(@NotNull PsiElement parent, @NotNull String referencedName) {
+  @NotNull
+  private static List<RatedResolveResult> resolveMemberFromReferenceTypeProviders(@NotNull PsiElement parent,
+                                                                                  @NotNull String referencedName) {
     final PyResolveContext resolveContext = PyResolveContext.defaultContext();
-    PyType refType = PyReferenceExpressionImpl.getReferenceTypeFromProviders(parent, resolveContext.getTypeEvalContext(), null);
+    final PyType refType = PyReferenceExpressionImpl.getReferenceTypeFromProviders(parent, resolveContext.getTypeEvalContext(), null);
     if (refType != null) {
       final List<? extends RatedResolveResult> result = refType.resolveMember(referencedName, null, AccessDirection.READ, resolveContext);
-      if (result != null && !result.isEmpty()) {
-        return result.get(0).getElement();
+      if (result != null) {
+        return Lists.newArrayList(result);
       }
     }
-    return null;
+    return Collections.emptyList();
   }
 
   private static boolean isDunderAll(@NotNull PsiElement element) {
     return (element instanceof PyElement) && PyNames.ALL.equals(((PyElement)element).getName());
   }
 
-  @Nullable
-  private static PsiElement resolveInDirectory(final String referencedName, @Nullable final PsiFile containingFile,
-                                               final PsiDirectory dir, boolean isFileOnly, boolean checkForPackage) {
-    if (referencedName == null) return null;
-
+  @NotNull
+  private static List<RatedResolveResult> resolveInDirectory(@NotNull final String referencedName, @Nullable final PsiFile containingFile,
+                                                             final PsiDirectory dir, boolean isFileOnly, boolean checkForPackage) {
     final PsiDirectory subdir = dir.findSubdirectory(referencedName);
     if (subdir != null && (!checkForPackage || PyUtil.isPackage(subdir, containingFile))) {
-      return subdir;
+      return ResolveResultList.to(subdir);
     }
 
     final PsiFile module = findPyFileInDir(dir, referencedName);
-    if (module != null) return module;
+    if (module != null) {
+      return ResolveResultList.to(module);
+    }
 
     if (!isFileOnly) {
       // not a subdir, not a file; could be a name in parent/__init__.py
       final PsiFile initPy = dir.findFile(PyNames.INIT_DOT_PY);
-      if (initPy == containingFile) return null; // don't dive into the file we're in
+      if (initPy == containingFile) {
+        return Collections.emptyList(); // don't dive into the file we're in
+      }
       if (initPy instanceof PyFile) {
-        return ((PyFile)initPy).getElementNamed(referencedName);
+        return ((PyFile)initPy).multiResolveName(referencedName);
       }
     }
-    return null;
+    return Collections.emptyList();
   }
 
   @Nullable

python/src/com/jetbrains/python/psi/types/PyModuleType.java

@@ -86,9 +86,9 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
     if (overridingMember != null) {
       return ResolveResultList.to(overridingMember);
     }
-    final PsiElement attribute = myModule.getElementNamed(name);
-    if (attribute != null) {
-      return ResolveResultList.to(attribute);
+    final List<RatedResolveResult> attributes = myModule.multiResolveName(name);
+    if (!attributes.isEmpty()) {
+      return attributes;
     }
     if (PyUtil.isPackage(myModule)) {
       final List<PyImportElement> importElements = new ArrayList<PyImportElement>();

python/src/com/jetbrains/python/refactoring/PyDefUseUtil.java

@@ -84,7 +84,7 @@ public class PyDefUseUtil {
                                       }
                                     }
                                     else if (acceptImplicitImports && implicit != null) {
-                                      if (implicit.getElementNamed(varName) != null) {
+                                      if (!implicit.multiResolveName(varName).isEmpty()) {
                                         result.add(instruction);
                                         return ControlFlowUtil.Operation.CONTINUE;
                                       }

python/testSrc/com/jetbrains/python/PyTypeTest.java

@@ -1004,6 +1004,7 @@ public class PyTypeTest extends PyTestCase {
            "    expr = foo\n");
   }
 
+  // PY-18217
   public void testConditionImportOuterScope() {
     doMultiFileTest("Union[str, int]",
                     "if something:\n" +
@@ -1015,8 +1016,9 @@ public class PyTypeTest extends PyTestCase {
                     "    expr = foo\n");
   }
 
+  // PY-18402
   public void testConditionInImportedModule() {
-    doMultiFileTest("Union[str, int]",
+    doMultiFileTest("Union[int, str]",
                     "from m1 import foo\n" +
                     "\n" +
                     "def f():\n" +