PY-18816 Resolve generic classes and type parameters using stubs

This commit is contained in:
Mikhail Golubev
2017-07-19 19:28:31 +03:00
parent 80fee7c7bb
commit 2cc6c4daf0
5 changed files with 144 additions and 42 deletions
@@ -19,7 +19,6 @@ import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.intellij.openapi.extensions.Extensions;
import com.intellij.openapi.util.Ref;
import com.intellij.psi.PsiElement;
import com.intellij.psi.impl.source.resolve.FileContextUtil;
@@ -39,8 +38,12 @@ import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyBuiltinCache;
import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
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.*;
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.stubs.PyClassStub;
import com.jetbrains.python.psi.types.*;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NotNull;
@@ -353,35 +356,78 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
return Collections.emptyMap();
}
final Map<PyType, PyType> results = new HashMap<>();
// XXX: Requires switching from stub to AST
for (PyExpression e : cls.getSuperClassExpressions()) {
final PySubscriptionExpression subscriptionExpr = as(e, PySubscriptionExpression.class);
final PyExpression superExpr = subscriptionExpr != null ? subscriptionExpr.getOperand() : e;
final PyType superType = context.getType(superExpr);
final PyClassType superClassType = as(superType, PyClassType.class);
final PyClass superClass = superClassType != null ? superClassType.getPyClass() : null;
final Map<PyType, PyType> superSubstitutions = superClass != null
? doPreventingRecursion(RECURSION_KEY, false,
() -> getGenericSubstitutions(superClass, context))
: null;
for (Map.Entry<PyClass, PySubscriptionExpression> e : getResolvedSuperClassesAndTypeParameters(cls, context).entrySet()) {
final PySubscriptionExpression subscriptionExpr = e.getValue();
final PyClass superClass = e.getKey();
final Map<PyType, PyType> superSubstitutions = doPreventingRecursion(RECURSION_KEY, false, () -> getGenericSubstitutions(superClass, context));
if (superSubstitutions != null) {
results.putAll(superSubstitutions);
}
final List<PyType> superGenerics = superClass != null ? collectGenericTypes(superClass, ctx) : Collections.emptyList();
final List<PyExpression> indices = subscriptionExpr != null ? getSubscriptionIndices(subscriptionExpr) : Collections.emptyList();
for (int i = 0; i < superGenerics.size(); i++) {
final PyExpression expr = ContainerUtil.getOrElse(indices, i, null);
final PyType superGeneric = superGenerics.get(i);
final Ref<PyType> typeRef = expr != null ? getType(expr, ctx) : null;
final PyType actualType = typeRef != null ? typeRef.get() : null;
if (!superGeneric.equals(actualType)) {
results.put(superGeneric, actualType);
if (superClass != null) {
final List<PyType> superGenerics = collectGenericTypes(superClass, ctx);
final List<PyExpression> indices = subscriptionExpr != null ? getSubscriptionIndices(subscriptionExpr) : Collections.emptyList();
for (int i = 0; i < superGenerics.size(); i++) {
final PyExpression expr = ContainerUtil.getOrElse(indices, i, null);
final PyType superGeneric = superGenerics.get(i);
final Ref<PyType> typeRef = expr != null ? getType(expr, ctx) : null;
final PyType actualType = typeRef != null ? typeRef.get() : null;
if (!superGeneric.equals(actualType)) {
results.put(superGeneric, actualType);
}
}
}
}
return results;
}
@NotNull
private static Map<PyClass, PySubscriptionExpression> getResolvedSuperClassesAndTypeParameters(@NotNull PyClass pyClass,
@NotNull TypeEvalContext context) {
final Map<PyClass, PySubscriptionExpression> results = new LinkedHashMap<>();
final PyClassStub classStub = pyClass.getStub();
if (context.maySwitchToAST(pyClass)) {
for (PyExpression e : pyClass.getSuperClassExpressions()) {
final PySubscriptionExpression subscriptionExpr = as(e, PySubscriptionExpression.class);
final PyExpression superExpr = subscriptionExpr != null ? subscriptionExpr.getOperand() : e;
final PyType superType = context.getType(superExpr);
final PyClassType superClassType = as(superType, PyClassType.class);
final PyClass superClass = superClassType != null ? superClassType.getPyClass() : null;
if (superClass != null) {
results.put(superClass, subscriptionExpr);
}
}
return results;
}
final Iterable<QualifiedName> allBaseClassesQNames;
final List<PySubscriptionExpression> subscriptedBaseClasses = PyClassElementType.getSubscriptedSuperClassesStubSafe(pyClass);
final Map<QualifiedName, PySubscriptionExpression> baseClassQNameToExpr = new HashMap<>();
if (classStub == null) {
allBaseClassesQNames = PyClassElementType.getSuperClassQNames(pyClass).keySet();
}
else {
allBaseClassesQNames = classStub.getSuperClasses().keySet();
}
for (PySubscriptionExpression subscriptedBase : subscriptedBaseClasses) {
final PyExpression operand = subscriptedBase.getOperand();
if (operand instanceof PyReferenceExpression) {
final QualifiedName className = PyPsiUtils.asQualifiedName(operand);
baseClassQNameToExpr.put(className, subscriptedBase);
}
}
for (QualifiedName qName : allBaseClassesQNames) {
final List<PsiElement> classes = 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) {
results.put(firstFound, baseClassQNameToExpr.get(qName));
}
}
return results;
}
@NotNull
private static List<PyExpression> getSubscriptionIndices(@NotNull PySubscriptionExpression expr) {
final PyExpression indexExpr = expr.getIndexExpression();
@@ -396,8 +442,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
}
final TypeEvalContext typeEvalContext = context.getTypeContext();
// XXX: Requires switching from stub to AST
return StreamEx.of(cls.getSuperClassExpressions())
.select(PySubscriptionExpression.class)
return StreamEx.of(PyClassElementType.getSubscriptedSuperClassesStubSafe(cls))
.map(PySubscriptionExpression::getIndexExpression)
.flatMap(e -> {
final PyTupleExpression tupleExpr = as(e, PyTupleExpression.class);
@@ -786,7 +831,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
if (element instanceof PyTargetExpression) {
final PyTargetExpression targetExpr = (PyTargetExpression)element;
final PyExpression assignedValue;
if (context.maySwitchToAST(expression)) {
if (context.maySwitchToAST(targetExpr)) {
assignedValue = targetExpr.findAssignedValue();
}
else {
@@ -823,16 +868,25 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
final QualifiedName qualifiedName = turnPlainReferenceExpressionIntoQualifiedName(expression);
final PyFile pyFile = as(FileContextUtil.getContextFile(expression), PyFile.class);
if (pyFile != null && qualifiedName != null && qualifiedName.getComponentCount() > 0) {
if (pyFile != null && qualifiedName != null) {
return 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(qualifiedName.getFirstComponent(), false));
results.addAll(pyFile.multiResolveName(first, false));
if (results.isEmpty()) {
for (PyReferenceResolveProvider provider : Extensions.getExtensions(PyReferenceResolveProvider.EP_NAME)) {
if (provider instanceof PyOverridingReferenceResolveProvider) {
continue;
}
results.addAll(provider.resolveName(expression, context));
final PsiElement builtinSymbol = PyBuiltinCache.getInstance(pyFile).getByName(first);
if (builtinSymbol != null) {
results.add(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, builtinSymbol));
}
}
@@ -846,7 +900,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
final PyType type = context.getType((PyTypedElement)element);
if (type != null) {
final List<? extends RatedResolveResult> resolved =
type.resolveMember(name, expression, AccessDirection.READ, resolveContext);
type.resolveMember(name, null, AccessDirection.READ, resolveContext);
if (resolved != null) {
children.addAll(resolved);
}
@@ -857,7 +911,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
}
return PyUtil.filterTopPriorityResults(results.toArray(RatedResolveResult.EMPTY_ARRAY));
}
return Collections.singletonList(expression);
return Collections.emptyList();
}
/**
@@ -19,14 +19,15 @@ import com.intellij.lang.ASTNode;
import com.intellij.psi.PsiElement;
import com.intellij.psi.stubs.*;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.intellij.util.io.StringRef;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyClassImpl;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import com.jetbrains.python.psi.resolve.PyResolveUtil;
import com.jetbrains.python.psi.stubs.*;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
@@ -34,6 +35,8 @@ import org.jetbrains.annotations.Nullable;
import java.io.IOException;
import java.util.*;
import static com.jetbrains.python.psi.PyUtil.as;
/**
* @author max
*/
@@ -61,7 +64,7 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
return new PyClassStubImpl(psi.getName(),
parentStub,
getSuperClassQNames(psi),
getSubscriptedSuperClasses(psi),
ContainerUtil.map(getSubscriptedSuperClasses(psi), PsiElement::getText),
PyPsiUtils.asQualifiedName(psi.getMetaClassExpression()),
psi.getOwnSlots(),
PyPsiUtils.strValue(psi.getDocStringExpression()),
@@ -83,11 +86,18 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
}
@NotNull
private static List<String> getSubscriptedSuperClasses(@NotNull PyClass pyClass) {
return StreamEx.of(pyClass.getSuperClassExpressions())
.filter(PySubscriptionExpression.class::isInstance)
.map(PsiElement::getText)
.toList();
private static List<PySubscriptionExpression> getSubscriptedSuperClasses(@NotNull PyClass pyClass) {
return ContainerUtil.mapNotNull(pyClass.getSuperClassExpressions(), x -> as(x, PySubscriptionExpression.class));
}
@NotNull
public static List<PySubscriptionExpression> getSubscriptedSuperClassesStubSafe(@NotNull PyClass pyClass) {
final PyClassStub classStub = pyClass.getStub();
if (classStub == null) {
return getSubscriptedSuperClasses(pyClass);
}
return ContainerUtil.mapNotNull(classStub.getSubscriptedSuperClasses(),
x -> as(PyTypingTypeProvider.createExpressionFromFragment(x, pyClass), PySubscriptionExpression.class));
}
@Nullable
@@ -0,0 +1,11 @@
from mod import Base, T1
class MyClass(Base[T1, None]):
pass
x = MyClass(42)
expr = x.m('foo')
print(expr)
@@ -0,0 +1,13 @@
from typing import TypeVar, Generic, Tuple
T1 = TypeVar('T1')
T2 = TypeVar('T2')
T3 = TypeVar('T3')
class Base(Generic[T1, T2, T3]):
def __init__(self, x: T1):
pass
def m(self, x: T3) -> Tuple[T1, T2, T3]:
pass
@@ -750,7 +750,7 @@ public class PyStubsTest extends PyTestCase {
assertNotParsed(file);
});
}
// PY-18116
public void testVariableAnnotation() {
runWithLanguageLevel(LanguageLevel.PYTHON36, () -> {
@@ -835,4 +835,18 @@ public class PyStubsTest extends PyTestCase {
assertContainsOrdered(genericBases, "Generic[T, V]");
assertNotParsed(file);
}
// PY-18816
public void testComplexGenericType() {
runWithLanguageLevel(LanguageLevel.PYTHON30, () -> {
myFixture.copyDirectoryToProject(getTestName(true), "");
final PsiManager manager = PsiManager.getInstance(myFixture.getProject());
final PyFile originFile = (PyFile)manager.findFile(myFixture.findFileInTempDir("a.py"));
final PyFile libFile = (PyFile)manager.findFile(myFixture.findFileInTempDir("mod.py"));
final PyTargetExpression instance = originFile.findTopLevelAttribute("expr");
assertType("Tuple[int, None, str]", instance, TypeEvalContext.codeAnalysis(myFixture.getProject(), originFile));
assertNotParsed(libFile);
});
}
}