From 2cc6c4daf07330c817a95b0203dc2e7def7ad259 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 29 Jun 2017 15:30:20 +0300 Subject: [PATCH] PY-18816 Resolve generic classes and type parameters using stubs --- .../typing/PyTypingTypeProvider.java | 122 +++++++++++++----- .../psi/impl/stubs/PyClassElementType.java | 24 +++- python/testData/stubs/complexGenericType/a.py | 11 ++ .../testData/stubs/complexGenericType/mod.py | 13 ++ .../com/jetbrains/python/PyStubsTest.java | 16 ++- 5 files changed, 144 insertions(+), 42 deletions(-) create mode 100644 python/testData/stubs/complexGenericType/a.py create mode 100644 python/testData/stubs/complexGenericType/mod.py diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index 5dc7e5f700a9..1c3e7fd204a2 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -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 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 superSubstitutions = superClass != null - ? doPreventingRecursion(RECURSION_KEY, false, - () -> getGenericSubstitutions(superClass, context)) - : null; + + for (Map.Entry e : getResolvedSuperClassesAndTypeParameters(cls, context).entrySet()) { + final PySubscriptionExpression subscriptionExpr = e.getValue(); + final PyClass superClass = e.getKey(); + final Map superSubstitutions = doPreventingRecursion(RECURSION_KEY, false, () -> getGenericSubstitutions(superClass, context)); if (superSubstitutions != null) { results.putAll(superSubstitutions); } - final List superGenerics = superClass != null ? collectGenericTypes(superClass, ctx) : Collections.emptyList(); - final List 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 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 superGenerics = collectGenericTypes(superClass, ctx); + final List 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 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 getResolvedSuperClassesAndTypeParameters(@NotNull PyClass pyClass, + @NotNull TypeEvalContext context) { + final Map 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 allBaseClassesQNames; + final List subscriptedBaseClasses = PyClassElementType.getSubscriptedSuperClassesStubSafe(pyClass); + final Map 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 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 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 resolveQualifiedNameInFile(@NotNull QualifiedName qualifiedName, + @NotNull PyFile pyFile, + @NotNull TypeEvalContext context) { + if (qualifiedName.getComponentCount() > 0) { List 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 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(); } /** diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java index ce533297f659..232d7076ea15 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java @@ -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 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 } @NotNull - private static List getSubscriptedSuperClasses(@NotNull PyClass pyClass) { - return StreamEx.of(pyClass.getSuperClassExpressions()) - .filter(PySubscriptionExpression.class::isInstance) - .map(PsiElement::getText) - .toList(); + private static List getSubscriptedSuperClasses(@NotNull PyClass pyClass) { + return ContainerUtil.mapNotNull(pyClass.getSuperClassExpressions(), x -> as(x, PySubscriptionExpression.class)); + } + + @NotNull + public static List 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 diff --git a/python/testData/stubs/complexGenericType/a.py b/python/testData/stubs/complexGenericType/a.py new file mode 100644 index 000000000000..278015dec04f --- /dev/null +++ b/python/testData/stubs/complexGenericType/a.py @@ -0,0 +1,11 @@ +from mod import Base, T1 + + +class MyClass(Base[T1, None]): + pass + + +x = MyClass(42) +expr = x.m('foo') + +print(expr) diff --git a/python/testData/stubs/complexGenericType/mod.py b/python/testData/stubs/complexGenericType/mod.py new file mode 100644 index 000000000000..96643f7c421d --- /dev/null +++ b/python/testData/stubs/complexGenericType/mod.py @@ -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 \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java index d9c378b30087..c407622400c6 100644 --- a/python/testSrc/com/jetbrains/python/PyStubsTest.java +++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java @@ -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); + }); + } }