PY-18816 Extract common algorithm for resolving symbols by their qualified names

Use it both to resolve base and metaclasses by their qualified names
saved in stubs in PyClassImpl and to process type annotations extracted
from stubs in PyTypingTypeProvider.

python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -40,9 +40,8 @@ import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.stubs.PyClassElementType;
 import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType;
-import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.PyResolveImportUtil;
-import com.jetbrains.python.psi.resolve.RatedResolveResult;
+import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.stubs.PyClassStub;
 import com.jetbrains.python.psi.types.*;
 import one.util.streamex.StreamEx;
@@ -418,7 +417,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
       }
     }
     for (QualifiedName qName : allBaseClassesQNames) {
-      final List<PsiElement> classes = resolveQualifiedNameInFile(qName, (PyFile)pyClass.getContainingFile(), context);
+      final List<PsiElement> classes = PyResolveUtil.resolveQualifiedNameInFile(qName, (PyFile)pyClass.getContainingFile(), context);
       // Better way to handle results of the multiresove
       final PyClass firstFound = ContainerUtil.findInstance(classes, PyClass.class);
       if (firstFound != null) {
@@ -865,51 +864,11 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     final PyFile pyFile = as(FileContextUtil.getContextFile(expression), PyFile.class);
 
     if (pyFile != null && qualifiedName != null) {
-      return resolveQualifiedNameInFile(qualifiedName, pyFile, context);
+      return PyResolveUtil.resolveQualifiedNameInFile(qualifiedName, pyFile, context);
     }
     return Collections.singletonList(expression);
   }
 
-  @NotNull
-  private static List<PsiElement> resolveQualifiedNameInFile(@NotNull QualifiedName qualifiedName,
-                                                             @NotNull PyFile pyFile,
-                                                             @NotNull TypeEvalContext context) {
-    if (qualifiedName.getComponentCount() > 0) {
-      List<RatedResolveResult> results = new ArrayList<>();
-      final String first = qualifiedName.getFirstComponent();
-      //noinspection ConstantConditions
-      results.addAll(pyFile.multiResolveName(first, false));
-      if (results.isEmpty()) {
-        final PsiElement builtinSymbol = PyBuiltinCache.getInstance(pyFile).getByName(first);
-        if (builtinSymbol != null) {
-          results.add(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, builtinSymbol));
-        }
-      }
-
-      final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
-
-      for (String name : qualifiedName.removeHead(1).getComponents()) {
-        final List<RatedResolveResult> children = new ArrayList<>();
-        for (RatedResolveResult result : results) {
-          final PsiElement element = result.getElement();
-          if (element instanceof PyTypedElement) {
-            final PyType type = context.getType((PyTypedElement)element);
-            if (type != null) {
-              final List<? extends RatedResolveResult> resolved =
-                type.resolveMember(name, null, AccessDirection.READ, resolveContext);
-              if (resolved != null) {
-                children.addAll(resolved);
-              }
-            }
-          }
-        }
-        results = children;
-      }
-      return PyUtil.filterTopPriorityResults(results.toArray(RatedResolveResult.EMPTY_ARRAY));
-    }
-    return Collections.emptyList();
-  }
-
   /**
    * Return the qualified name containing all names in the given (possibly qualified) reference expression.
    * If any of the qualifiers is not a reference expression, returns null.
@@ -917,7 +876,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
    * For instance, for the expression "foo.bar.baz" it returns the qualified name "foo.bar.baz",
    * but for "foo[0].bar.baz" it will return null.
    * <p>
-   * If you need to take into account such implicit "magical" names, use {@link com.jetbrains.python.psi.impl.PyPsiUtils#asQualifiedName(PyExpression)}
+   * If you need to take into account such implicit "magical" names, use {@link PyPsiUtils#asQualifiedName(PyExpression)}
    * or {@link PyQualifiedExpression#asQualifiedName()}.
    * @param expression
    */

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

@@ -44,7 +44,6 @@ import com.jetbrains.python.psi.impl.stubs.PyClassElementType;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
-import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.stubs.PropertyStubStorage;
 import com.jetbrains.python.psi.stubs.PyClassStub;
 import com.jetbrains.python.psi.stubs.PyFunctionStub;
@@ -1708,73 +1707,10 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
     return results;
   }
 
-  @Nullable
-  private static PsiElement getElementQNamed(@NotNull PyFile file, @NotNull QualifiedName qualifiedName, @NotNull TypeEvalContext context) {
-    if (qualifiedName.getComponentCount() == 0) {
-      return null;
-    }
-    final String first = qualifiedName.getFirstComponent();
-    if (first == null) {
-      return null;
-    }
-    final QualifiedName rest = qualifiedName.removeHead(1);
-    final PsiElement firstElement = file.multiResolveName(first, false)
-      .stream()
-      .map(RatedResolveResult::getElement)
-      .findFirst()
-      .orElse(PyBuiltinCache.getInstance(file).getByName(first));
-    if (rest.getComponentCount() == 0) {
-      return firstElement;
-    }
-    final PyTypedElement typedElement = as(firstElement, PyTypedElement.class);
-    if (typedElement == null) {
-      return null;
-    }
-    final String name = rest.getLastComponent();
-    final QualifiedName containingQName = rest.removeLastComponent();
-    PyType currentType = context.getType(typedElement);
-    if (currentType == null) {
-      return null;
-    }
-    for (String component : containingQName.getComponents()) {
-      currentType = getMemberType(currentType, component, context);
-      if (currentType == null) {
-        return null;
-      }
-    }
-    if (name != null) {
-      return resolveTypeMember(currentType, name, context);
-    }
-    return null;
-  }
-
-  @Nullable
-  private static PyType getMemberType(@NotNull PyType type, @NotNull String name, @NotNull TypeEvalContext context) {
-    final PyType result;
-    PsiElement element = resolveTypeMember(type, name, context);
-    if (element instanceof PyTypedElement) {
-      result = context.getType((PyTypedElement)element);
-    }
-    else {
-      return null;
-    }
-    if (result instanceof PyClassLikeType) {
-      return ((PyClassLikeType)result).toInstance();
-    }
-    return result;
-  }
-
-  @Nullable
-  private static PsiElement resolveTypeMember(@NotNull PyType type, @NotNull String name, @NotNull TypeEvalContext context) {
-    final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
-    final List<? extends RatedResolveResult> results = type.resolveMember(name, null, AccessDirection.READ, resolveContext);
-    return (results != null && !results.isEmpty()) ? results.get(0).getElement() : null;
-  }
-
   @Nullable
   private static PyClassLikeType classTypeFromQName(@NotNull QualifiedName qualifiedName, @NotNull PyFile containingFile,
                                                     @NotNull TypeEvalContext context) {
-    final PsiElement element = getElementQNamed(containingFile, qualifiedName, context);
+    final PsiElement element = ContainerUtil.getFirstItem(PyResolveUtil.resolveQualifiedNameInFile(qualifiedName, containingFile, context));
     if (element instanceof PyTypedElement) {
       final PyType type = context.getType((PyTypedElement)element);
       if (type instanceof PyClassLikeType) {

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

@@ -19,17 +19,25 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.ResolveState;
 import com.intellij.psi.scope.PsiScopeProcessor;
+import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
+
+import static com.jetbrains.python.psi.PyUtil.as;
 
 /**
  * @author vlan
@@ -125,4 +133,50 @@ public class PyResolveUtil {
 
     return processor.getElements();
   }
+
+  /**
+   * Resolve a symbol by its qualified name, starting from the specified file and then following the chain of type members.
+   * This type of resolve is stub-safe, i.e. it's not supposed to cause any un-stubbing of external files unless it explicitly
+   * allowed by the given type evaluation context.
+   *
+   * @param qualifiedName name of a symbol to resolve
+   * @param file          module which serves as a starting point for resolve
+   * @return all possible candidates that can be found by the given qualified name
+   */
+  @NotNull
+  public static List<PsiElement> resolveQualifiedNameInFile(@NotNull QualifiedName qualifiedName,
+                                                            @NotNull PyFile file,
+                                                            @NotNull TypeEvalContext context) {
+    final String firstName = qualifiedName.getFirstComponent();
+    if (firstName == null) {
+      return Collections.emptyList();
+    }
+    final List<RatedResolveResult> unqualifiedResults = file.multiResolveName(firstName, false);
+    final StreamEx<RatedResolveResult> initialResults;
+    if (unqualifiedResults.isEmpty()) {
+      final PsiElement builtin = PyBuiltinCache.getInstance(file).getByName(firstName);
+      if (builtin == null) {
+        return Collections.emptyList();
+      }
+      initialResults = StreamEx.of(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, builtin));
+    }
+    else {
+      initialResults = StreamEx.of(unqualifiedResults);
+    }
+    final List<String> remainingNames = qualifiedName.removeHead(1).getComponents();
+    final StreamEx<RatedResolveResult> result = StreamEx.of(remainingNames).reduce(initialResults, (prev, name) ->
+      prev
+        .map(RatedResolveResult::getElement)
+        .map(elem -> as(elem, PyTypedElement.class))
+        .nonNull()
+        .map(context::getType)
+        .nonNull()
+        .map(type -> type instanceof PyClassLikeType ? ((PyClassLikeType)type).toInstance() : type)
+        .flatMap(elemType -> {
+          final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
+          final List<? extends RatedResolveResult> results = elemType.resolveMember(name, null, AccessDirection.READ, resolveContext);
+          return results != null ? StreamEx.of(results) : StreamEx.<RatedResolveResult>empty();
+        }), (oldResults, newResults) -> newResults);
+    return PyUtil.filterTopPriorityResults(result.toArray(RatedResolveResult[]::new));
+  }
 }