Provide callable type for dataclass initialization (PY-27398)

Add `dataclasses` to stdlib_packages.txt

python/helpers/tools/stdlib_packages.txt

@@ -55,6 +55,7 @@ csv
 ctypes
 curses
 datetime
+dataclasses
 dbhash
 dbm
 decimal

python/src/META-INF/python-core-common.xml

@@ -703,6 +703,7 @@
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>
     <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider"/>
     <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibOverridingTypeProvider"/>
+    <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyDataclassesTypeProvider"/>
     <pyModuleMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibModuleMembersProvider"/>
     <documentationLinkProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibDocumentationLinkProvider"/>
     <canonicalPathProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibCanonicalPathProvider"/>

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+package com.jetbrains.python.codeInsight.stdlib
+
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyCallExpressionNavigator
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.*
+
+class PyDataclassesTypeProvider : PyTypeProviderBase() {
+
+  override fun getReferenceExpressionType(referenceExpression: PyReferenceExpression, context: TypeEvalContext): PyType? {
+    return getDataclassTypeForCallee(referenceExpression, context)
+  }
+
+  private fun getDataclassTypeForCallee(referenceExpression: PyReferenceExpression, context: TypeEvalContext): PyCallableType? {
+    if (PyCallExpressionNavigator.getPyCallExpressionByCallee(referenceExpression) == null) return null
+
+    val resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context)
+    val resolveResults = referenceExpression.getReference(resolveContext).multiResolve(false)
+
+    return PyUtil.filterTopPriorityResults(resolveResults)
+      .asSequence()
+      .filterIsInstance<PyClass>()
+      .map { getDataclassTypeForClass(it, context) }
+      .firstOrNull { it != null }
+  }
+
+  private fun getDataclassTypeForClass(cls: PyClass, context: TypeEvalContext): PyCallableType? {
+    if (!PyKnownDecoratorUtil.getKnownDecorators(cls, context).contains(PyKnownDecoratorUtil.KnownDecorator.DATACLASSES_DATACLASS)) {
+      return null
+    }
+
+    val parameters = ArrayList<PyCallableParameter>()
+
+    cls.processClassLevelDeclarations { element, _ ->
+      if (element is PyTargetExpression && element.annotationValue != null) {
+        parameters.add(PyCallableParameterImpl.nonPsi(element.name, context.getType(element), element.findAssignedValue()))
+      }
+      true
+    }
+
+    return PyCallableTypeImpl(parameters, context.getType(cls))
+  }
+}

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

@@ -67,7 +67,9 @@ public class PyKnownDecoratorUtil {
 
     PYRAMID_DECORATOR_REIFY("pyramid.decorator.reify"),
     DJANGO_UTILS_FUNCTIONAL_CACHED_PROPERTY("django.utils.functional.cached_property"),
-    KOMBU_UTILS_CACHED_PROPERTY("kombu.utils.cached_property");
+    KOMBU_UTILS_CACHED_PROPERTY("kombu.utils.cached_property"),
+
+    DATACLASSES_DATACLASS("dataclasses.dataclass");
 
     private final QualifiedName myQualifiedName;
 

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

@@ -0,0 +1,13 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A:
+    x: int
+    y: str
+    z: float = 0.0
+
+A(<warning descr="Parameter 'x' unfilled"><warning descr="Parameter 'y' unfilled">)</warning></warning>
+A(1<warning descr="Parameter 'y' unfilled">)</warning>
+A(1, "a")
+A(1, "a", 1.0)
+A(1, "a", 1.0, <warning descr="Unexpected argument">"b"</warning>)

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

@@ -0,0 +1,3 @@
+def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
+              hash=None, frozen=False):
+    pass

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

@@ -0,0 +1,13 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A:
+    x: int
+    y: str
+    z: float = 0.0
+
+A(1, "a")
+A(<warning descr="Expected type 'int', got 'str' instead">"a"</warning>, <warning descr="Expected type 'str', got 'int' instead">1</warning>)
+
+A(1, "a", 1.0)
+A(<warning descr="Expected type 'int', got 'str' instead">"a"</warning>, <warning descr="Expected type 'str', got 'int' instead">1</warning>, <warning descr="Expected type 'float', got 'str' instead">"b"</warning>)

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

@@ -0,0 +1,3 @@
+def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
+              hash=None, frozen=False):
+    pass

python/testData/paramInfo/InitializingDataclass/a.py

@@ -0,0 +1,9 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A:
+    x: int
+    y: str
+    z: float = 0.0
+
+A(<arg1>)

python/testData/paramInfo/InitializingDataclass/dataclasses.py

@@ -0,0 +1,3 @@
+def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
+              hash=None, frozen=False):
+    pass

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

@@ -707,6 +707,18 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
     );
   }
 
+  // PY-27398
+  public void testInitializingDataclass() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () -> {
+        final Map<String, PsiElement> marks = loadMultiFileTest(1);
+
+        feignCtrlP(marks.get("<arg1>").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "});
+      }
+    );
+  }
+
   /**
    * Imitates pressing of Ctrl+P; fails if results are not as expected.
    * @param offset offset of 'cursor' where Ctrl+P is pressed.

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

@@ -276,4 +276,9 @@ public class Py3TypeCheckerInspectionTest extends PyInspectionTestCase {
   public void testChainedComparisonsGenericMatching() {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
   }
+
+  // PY-27398
+  public void testInitializingDataclass() {
+    runWithLanguageLevel(LanguageLevel.PYTHON37, () -> super.doMultiFileTest());
+  }
 }

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

@@ -318,4 +318,9 @@ public class PyArgumentListInspectionTest extends PyInspectionTestCase {
   public void testAbstractMethod() {
     runWithLanguageLevel(LanguageLevel.PYTHON30, this::doTest);
   }
+
+  // PY-27398
+  public void testInitializingDataclass() {
+    runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest);
+  }
 }