IDEA-CR-55284: PY-38412 Fix unstubbing in case of TypedDict creation with class-based syntax

GitOrigin-RevId: 3d165abaa22d6fa31d63964c8a7fa828fdefd1dd

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypedDictTypeProvider.kt

@@ -143,21 +143,25 @@ class PyTypedDictTypeProvider : PyTypeProviderBase() {
         true
       }
 
-      val argumentList = cls.children.filterIsInstance<PyArgumentList>().firstOrNull()
-      val totalityValue = argumentList?.getKeywordArgument("total")
-      val fieldsRequired = if (totalityValue != null && totalityValue.valueExpression is PyBoolLiteralExpression)
-        (totalityValue.valueExpression as PyBoolLiteralExpression).value
-      else true
-
+      val totality = getTotality(cls)
       val toTDFields = Collectors.toMap<PyTargetExpression, String, PyTypedDictType.FieldTypeAndTotality, TDFields>(
         { it.name },
-        { field -> PyTypedDictType.FieldTypeAndTotality(context.getType(field), fieldsRequired) },
+        { field -> PyTypedDictType.FieldTypeAndTotality(context.getType(field), totality) },
         { _, v2 -> v2 },
         { TDFields() })
 
       return fields.stream().collect(toTDFields)
     }
 
+    private fun getTotality(cls: PyClass): Boolean {
+      return if (cls.stub != null) {
+        "total=False" !in cls.stub.superClassesText
+      }
+      else {
+        (cls.superClassExpressionList?.getKeywordArgument("total")?.valueExpression as? PyBoolLiteralExpression)?.value ?: true
+      }
+    }
+
     private fun getTypedDictTypeForTarget(target: PyTargetExpression, context: TypeEvalContext): PyTypedDictType? {
       val stub = target.stub
 
@@ -283,7 +287,8 @@ class PyTypedDictTypeProvider : PyTypeProviderBase() {
                                  PyTypedDictType.DefinitionLevel.INSTANCE,
                                  emptyList())
         }
-      } else if (expression is PyCallExpression) {
+      }
+      else if (expression is PyCallExpression) {
         val resolvedQualifiedNames = if (expression.callee != null) resolveToQualifiedNames(expression.callee!!, context) else return null
         if (resolvedQualifiedNames.any { it == PyNames.DICT }) {
           val arguments = expression.arguments

python/testData/inspections/PyTypeCheckerInspection/TypedDictInStub/a.py

@@ -0,0 +1,3 @@
+avrora = {'name': 'Iskra', 'id': 42}
+rodina = {'name': 'Iskra'}
+luksor = <warning descr="Expected type 'Cinema', got 'Dict[str, str]' instead">{'name': 'Iskra', 'id': '42'}</warning>

python/testData/inspections/PyTypeCheckerInspection/TypedDictInStub/a.pyi

@@ -0,0 +1,11 @@
+from typing import TypedDict
+
+
+class Cinema(TypedDict, total=False):
+    name: str
+    id: int
+
+
+avrora: Cinema
+rodina: Cinema
+luksor: Cinema

python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java

@@ -1083,4 +1083,9 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
       )
     );
   }
+
+  // PY-38412
+  public void testTypedDictInStub() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doMultiFileTest);
+  }
 }