diff --git a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java index 49935b62b20f..ad3242514157 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java +++ b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java @@ -36,6 +36,13 @@ public interface PyClassStub extends NamedStub { @NotNull Map getSuperClasses(); + /** + * Returns literal text of the subscription expressions in the base classes list. + * It's intended to be used for resolving generic types and type parameters in PEP 484 notation. + */ + @NotNull + List getSubscriptedSuperClasses(); + @Nullable QualifiedName getMetaClass(); diff --git a/python/src/com/jetbrains/python/psi/PyFileElementType.java b/python/src/com/jetbrains/python/psi/PyFileElementType.java index 484120a0d903..30e0d94372b7 100644 --- a/python/src/com/jetbrains/python/psi/PyFileElementType.java +++ b/python/src/com/jetbrains/python/psi/PyFileElementType.java @@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType { @Override public int getStubVersion() { // Don't forget to update versions of indexes that use the updated stub-based elements - return 62; + return 63; } @Nullable 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 15427b4386e8..ce533297f659 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java @@ -26,6 +26,7 @@ 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; @@ -60,6 +61,7 @@ public class PyClassElementType extends PyStubElementType return new PyClassStubImpl(psi.getName(), parentStub, getSuperClassQNames(psi), + getSubscriptedSuperClasses(psi), PyPsiUtils.asQualifiedName(psi.getMetaClassExpression()), psi.getOwnSlots(), PyPsiUtils.strValue(psi.getDocStringExpression()), @@ -80,6 +82,14 @@ public class PyClassElementType extends PyStubElementType return result; } + @NotNull + private static List getSubscriptedSuperClasses(@NotNull PyClass pyClass) { + return StreamEx.of(pyClass.getSuperClassExpressions()) + .filter(PySubscriptionExpression.class::isInstance) + .map(PsiElement::getText) + .toList(); + } + @Nullable private static QualifiedName resolveOriginalSuperClassQName(@NotNull PyExpression superClassExpression) { if (superClassExpression instanceof PyReferenceExpression) { @@ -116,6 +126,12 @@ public class PyClassElementType extends PyStubElementType QualifiedName.serialize(entry.getValue(), dataStream); } + final List baseClassesText = pyClassStub.getSubscriptedSuperClasses(); + dataStream.writeByte(baseClassesText.size()); + for (String text : baseClassesText) { + dataStream.writeName(text); + } + QualifiedName.serialize(pyClassStub.getMetaClass(), dataStream); PyFileElementType.writeNullableList(dataStream, pyClassStub.getSlots()); @@ -134,6 +150,15 @@ public class PyClassElementType extends PyStubElementType superClasses.put(QualifiedName.deserialize(dataStream), QualifiedName.deserialize(dataStream)); } + final byte parametrizedBaseClassesCount = dataStream.readByte(); + final ArrayList parametrizedBaseClasses = new ArrayList<>(); + for (int i = 0; i < parametrizedBaseClassesCount; i++) { + final StringRef ref = dataStream.readName(); + if (ref != null) { + parametrizedBaseClasses.add(ref.getString()); + } + } + final QualifiedName metaClass = QualifiedName.deserialize(dataStream); final List slots = PyFileElementType.readNullableList(dataStream); @@ -141,7 +166,7 @@ public class PyClassElementType extends PyStubElementType final String docStringInStub = dataStream.readUTFFast(); final String docString = docStringInStub.length() > 0 ? docStringInStub : null; - return new PyClassStubImpl(name, parentStub, superClasses, metaClass, slots, docString, getStubElementType()); + return new PyClassStubImpl(name, parentStub, superClasses, parametrizedBaseClasses, metaClass, slots, docString, getStubElementType()); } public void indexStub(@NotNull final PyClassStub stub, @NotNull final IndexSink sink) { diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java index 5a8e4d413dde..8afa13e1fb76 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java @@ -37,6 +37,7 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { @NotNull private final Map mySuperClasses; + private final List mySuperClassesText; @Nullable private final QualifiedName myMetaClass; @@ -50,6 +51,7 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { public PyClassStubImpl(@Nullable String name, @Nullable StubElement parentStub, @NotNull Map superClasses, + @NotNull List superClassesText, @Nullable QualifiedName metaClass, @Nullable List slots, @Nullable String docString, @@ -57,6 +59,7 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { super(parentStub, stubElementType); myName = name; mySuperClasses = superClasses; + mySuperClassesText = superClassesText; myMetaClass = metaClass; mySlots = slots; myDocString = docString; @@ -72,6 +75,12 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { return mySuperClasses; } + @NotNull + @Override + public List getSubscriptedSuperClasses() { + return mySuperClassesText; + } + @Nullable @Override public QualifiedName getMetaClass() { diff --git a/python/testData/stubs/ParametrizedBaseClass.py b/python/testData/stubs/ParametrizedBaseClass.py new file mode 100644 index 000000000000..5492737f5746 --- /dev/null +++ b/python/testData/stubs/ParametrizedBaseClass.py @@ -0,0 +1,8 @@ +from typing import Generic, TypeVar + +T = TypeVar('T') +V = TypeVar('V') + + +class Class(Generic[T, V]): + pass diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java index 60a1c931713f..d9c378b30087 100644 --- a/python/testSrc/com/jetbrains/python/PyStubsTest.java +++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java @@ -824,4 +824,15 @@ public class PyStubsTest extends PyTestCase { assertNull(message, stub); } } + + // PY-18816 + public void testParametrizedBaseClass() { + final PyFile file = getTestFile(); + final PyClass genericClass = file.findTopLevelClass("Class"); + final PyClassStub stub = genericClass.getStub(); + assertNotNull(stub); + final List genericBases = stub.getSubscriptedSuperClasses(); + assertContainsOrdered(genericBases, "Generic[T, V]"); + assertNotParsed(file); + } }