Correct types for initvars in __init__ (PY-27398)

python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt

@@ -4,7 +4,10 @@
 package com.jetbrains.python.codeInsight.stdlib
 
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
-import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyReferenceExpression
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.PyUtil
 import com.jetbrains.python.psi.impl.PyCallExpressionNavigator
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.types.*
@@ -41,7 +44,7 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
         val annotation = element.annotation
 
         if (annotation != null && !PyTypingTypeProvider.isClassVarAnnotation(annotation, context)) {
-          parameters.add(PyCallableParameterImpl.nonPsi(element.name, context.getType(element), element.findAssignedValue()))
+          parameters.add(PyCallableParameterImpl.nonPsi(element.name, getTypeForParameter(element, context), element.findAssignedValue()))
         }
       }
 
@@ -50,4 +53,12 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
 
     return PyCallableTypeImpl(parameters, context.getType(cls))
   }
+
+  private fun getTypeForParameter(element: PyTargetExpression, context: TypeEvalContext): PyType? {
+    val type = context.getType(element)
+    if (type is PyCollectionType && type is PyClassType && type.classQName == "dataclasses.InitVar") {
+      return type.elementTypes.firstOrNull()
+    }
+    return type
+  }
 }

python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py

@@ -74,3 +74,14 @@ class C2:
 C2(<warning descr="Parameter 'b' unfilled">)</warning>
 C2(1)
 C2(1, <warning descr="Unexpected argument">2</warning>)
+
+
+@dataclasses.dataclass
+class D1:
+    a: dataclasses.InitVar[int]
+    b: int
+
+D1(<warning descr="Parameter 'a' unfilled"><warning descr="Parameter 'b' unfilled">)</warning></warning>
+D1(1<warning descr="Parameter 'b' unfilled">)</warning>
+D1(1, 2)
+D1(1, 2, <warning descr="Unexpected argument">3</warning>)

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

@@ -73,4 +73,13 @@ class C2:
     b: int
 
 C2(1)
-C2(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>)
+C2(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>)
+
+
+@dataclasses.dataclass
+class D1:
+    a: dataclasses.InitVar[int]
+    b: int
+
+D1(1, 2)
+D1(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>, <warning descr="Expected type 'int', got 'str' instead">"2"</warning>)

python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py

@@ -1,3 +1,11 @@
+class _InitVarMeta(type):
+    def __getitem__(self, params):
+        return self
+
+class InitVar(metaclass=_InitVarMeta):
+    pass
+
+
 def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
               hash=None, frozen=False):
     pass

python/testData/paramInfo/InitializingDataclass/a.py

@@ -55,4 +55,12 @@ class C2:
     a: typing.ClassVar
     b: int
 
-C2(<arg6>)
+C2(<arg6>)
+
+
+@dataclasses.dataclass
+class D1:
+    a: dataclasses.InitVar[int]
+    b: int
+
+D1(<arg7>)

python/testData/paramInfo/InitializingDataclass/dataclasses.py

@@ -1,3 +1,11 @@
+class _InitVarMeta(type):
+    def __getitem__(self, params):
+        return self
+
+class InitVar(metaclass=_InitVarMeta):
+    pass
+
+
 def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
               hash=None, frozen=False):
     pass

python/testSrc/com/jetbrains/python/PyParameterInfoTest.java

@@ -712,7 +712,7 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
     runWithLanguageLevel(
       LanguageLevel.PYTHON37,
       () -> {
-        final Map<String, PsiElement> marks = loadMultiFileTest(6);
+        final Map<String, PsiElement> marks = loadMultiFileTest(7);
 
         feignCtrlP(marks.get("<arg1>").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "});
         feignCtrlP(marks.get("<arg2>").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "});
@@ -726,6 +726,7 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
         feignCtrlP(marks.get("<arg4>").getTextOffset()).check("self: B2, x: int", new String[]{"x: int"}, new String[]{"self: B2, "});
         feignCtrlP(marks.get("<arg5>").getTextOffset()).check("b: int", new String[]{"b: int"});
         feignCtrlP(marks.get("<arg6>").getTextOffset()).check("b: int", new String[]{"b: int"});
+        feignCtrlP(marks.get("<arg7>").getTextOffset()).check("a: int, b: int", new String[]{"a: int, "});
       }
     );
   }