PY-18816 Store literal text of subscripted base classes in PyClassStub

python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java

@@ -36,6 +36,13 @@ public interface PyClassStub extends NamedStub<PyClass> {
   @NotNull
   Map<QualifiedName, QualifiedName> 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<String> getSubscriptedSuperClasses();
+
   @Nullable
   QualifiedName getMetaClass();
 

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 62;
+    return 63;
   }
 
   @Nullable

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<PyClassStub, PyClass>
     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<PyClassStub, PyClass>
     return result;
   }
 
+  @NotNull
+  private static List<String> 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<PyClassStub, PyClass>
       QualifiedName.serialize(entry.getValue(), dataStream);
     }
 
+    final List<String> 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<PyClassStub, PyClass>
       superClasses.put(QualifiedName.deserialize(dataStream), QualifiedName.deserialize(dataStream));
     }
 
+    final byte parametrizedBaseClassesCount = dataStream.readByte();
+    final ArrayList<String> 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<String> slots = PyFileElementType.readNullableList(dataStream);
@@ -141,7 +166,7 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
     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) {

python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java

@@ -37,6 +37,7 @@ public class PyClassStubImpl extends StubBase<PyClass> implements PyClassStub {
 
   @NotNull
   private final Map<QualifiedName, QualifiedName> mySuperClasses;
+  private final List<String> mySuperClassesText;
 
   @Nullable
   private final QualifiedName myMetaClass;
@@ -50,6 +51,7 @@ public class PyClassStubImpl extends StubBase<PyClass> implements PyClassStub {
   public PyClassStubImpl(@Nullable String name,
                          @Nullable StubElement parentStub,
                          @NotNull Map<QualifiedName, QualifiedName> superClasses,
+                         @NotNull List<String> superClassesText,
                          @Nullable QualifiedName metaClass,
                          @Nullable List<String> slots,
                          @Nullable String docString,
@@ -57,6 +59,7 @@ public class PyClassStubImpl extends StubBase<PyClass> 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<PyClass> implements PyClassStub {
     return mySuperClasses;
   }
 
+  @NotNull
+  @Override
+  public List<String> getSubscriptedSuperClasses() {
+    return mySuperClassesText;
+  }
+
   @Nullable
   @Override
   public QualifiedName getMetaClass() {

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

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<String> genericBases = stub.getSubscriptedSuperClasses();
+    assertContainsOrdered(genericBases, "Generic[T, V]");
+    assertNotParsed(file);
+  }
 }