Support @dataclass(init=True/False) (PY-27398)

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

@@ -0,0 +1,115 @@
+/*
+ * 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.PyEvaluator
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+
+fun parseDataclassParameters(cls: PyClass, context: TypeEvalContext): DataclassParameters? {
+  val decorators = cls.decoratorList ?: return null
+
+  for (decorator in decorators.decorators) {
+    if (PyKnownDecoratorUtil.asKnownDecorators(decorator, context).contains(PyKnownDecoratorUtil.KnownDecorator.DATACLASSES_DATACLASS)) {
+      for (mapping in decorator.multiMapArguments(PyResolveContext.noImplicits().withTypeEvalContext(context))) {
+        if (mapping.unmappedArguments.isEmpty() && mapping.unmappedParameters.isEmpty()) {
+          val builder = DataclassParametersBuilder()
+
+          mapping
+            .mappedParameters
+            .entries
+            .forEach {
+              builder.update(it.value.name, it.key)
+            }
+
+          return builder.build()
+        }
+      }
+    }
+  }
+
+  return null
+}
+
+data class DataclassParameters(val init: Boolean,
+                               val repr: Boolean,
+                               val eq: Boolean,
+                               val order: Boolean,
+                               val hash: Boolean?,
+                               val frozen: Boolean,
+                               val initArgument: PyExpression?,
+                               val reprArgument: PyExpression?,
+                               val eqArgument: PyExpression?,
+                               val orderArgument: PyExpression?,
+                               val hashArgument: PyExpression?,
+                               val frozenArgument: PyExpression?)
+
+private class DataclassParametersBuilder {
+
+  companion object {
+    private val DEFAULT_INIT: Boolean = true
+    private val DEFAULT_REPR: Boolean = true
+    private val DEFAULT_EQ: Boolean = true
+    private val DEFAULT_ORDER: Boolean = false
+    private val DEFAULT_HASH: Boolean? = null // `null` means `None`
+    private val DEFAULT_FROZEN: Boolean = false
+  }
+
+  private var init = DEFAULT_INIT
+  private var repr = DEFAULT_REPR
+  private var eq = DEFAULT_EQ
+  private var order = DEFAULT_ORDER
+  private var hash = DEFAULT_HASH
+  private var frozen = DEFAULT_FROZEN
+
+  private var initArgument: PyExpression? = null
+  private var reprArgument: PyExpression? = null
+  private var eqArgument: PyExpression? = null
+  private var orderArgument: PyExpression? = null
+  private var hashArgument: PyExpression? = null
+  private var frozenArgument: PyExpression? = null
+
+  fun update(name: String?, argument: PyExpression?) {
+    val value = PyUtil.peelArgument(argument)
+
+    when (name) {
+      "init" -> {
+        init = evaluateBoolean(value, DEFAULT_INIT)
+        initArgument = argument
+      }
+      "repr" -> {
+        repr = evaluateBoolean(value, DEFAULT_REPR)
+        reprArgument = argument
+      }
+      "eq" -> {
+        eq = evaluateBoolean(value, DEFAULT_EQ)
+        eqArgument = argument
+      }
+      "order" -> {
+        order = evaluateBoolean(value, DEFAULT_ORDER)
+        orderArgument = argument
+      }
+      "hash" -> {
+        hash = evaluateBoolean(value, DEFAULT_HASH)
+        hashArgument = argument
+      }
+      "frozen" -> {
+        frozen = evaluateBoolean(value, DEFAULT_FROZEN)
+        frozenArgument = argument
+      }
+    }
+  }
+
+  fun build() = DataclassParameters(init, repr, eq, order, hash, frozen,
+                                    initArgument, reprArgument, eqArgument, orderArgument, hashArgument, frozenArgument)
+
+
+  private inline fun <reified T: Boolean?> evaluateBoolean(expression: PyExpression?, default: T): T {
+    val result = PyEvaluator().evaluate(expression)
+    if (result is Boolean) return result as T
+    return default
+  }
+}

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

@@ -28,7 +28,8 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
   }
 
   private fun getDataclassTypeForClass(cls: PyClass, context: TypeEvalContext): PyCallableType? {
-    if (!PyKnownDecoratorUtil.getKnownDecorators(cls, context).contains(PyKnownDecoratorUtil.KnownDecorator.DATACLASSES_DATACLASS)) {
+    val dataclassParameters = parseDataclassParameters(cls, context)
+    if (dataclassParameters == null || !dataclassParameters.init) {
       return null
     }
 

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

@@ -10,4 +10,46 @@ A(<warning descr="Parameter 'x' unfilled"><warning descr="Parameter 'y' unfilled
 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>)
+A(1, "a", 1.0, <warning descr="Unexpected argument">"b"</warning>)
+
+
+@dataclasses.dataclass(init=True)
+class A2:
+    x: int
+    y: str
+    z: float = 0.0
+
+A2(<warning descr="Parameter 'x' unfilled"><warning descr="Parameter 'y' unfilled">)</warning></warning>
+A2(1<warning descr="Parameter 'y' unfilled">)</warning>
+A2(1, "a")
+A2(1, "a", 1.0)
+A2(1, "a", 1.0, <warning descr="Unexpected argument">"b"</warning>)
+
+
+@dataclasses.dataclass(init=False)
+class B1:
+    x: int = 1
+    y: str = "2"
+    z: float = 0.0
+
+B1()
+B1<warning descr="Unexpected argument(s)Possible callees:object(self: object)object.__new__(cls: object)">(1)</warning>
+B1<warning descr="Unexpected argument(s)Possible callees:object(self: object)object.__new__(cls: object)">(1, "a")</warning>
+B1<warning descr="Unexpected argument(s)Possible callees:object(self: object)object.__new__(cls: object)">(1, "a", 1.0)</warning>
+B1<warning descr="Unexpected argument(s)Possible callees:object(self: object)object.__new__(cls: object)">(1, "a", 1.0, "b")</warning>
+
+
+@dataclasses.dataclass(init=False)
+class B2:
+    x: int
+    y: str
+    z: float = 0.0
+
+    def __init__(self, x: int):
+        self.x = x
+        self.y = str(x)
+        self.z = 0.0
+
+B2(<warning descr="Parameter 'x' unfilled">)</warning>
+B2(1)
+B2<warning descr="Unexpected argument(s)">(1, <warning descr="Unexpected argument">2</warning>)</warning>

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

@@ -10,4 +10,48 @@ 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>)
+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>)
+
+
+@dataclasses.dataclass(init=True)
+class A2:
+    x: int
+    y: str
+    z: float = 0.0
+
+A2(1, "a")
+A2(<warning descr="Expected type 'int', got 'str' instead">"a"</warning>, <warning descr="Expected type 'str', got 'int' instead">1</warning>)
+
+A2(1, "a", 1.0)
+A2(<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>)
+
+
+@dataclasses.dataclass(init=False)
+class B1:
+    x: int = 1
+    y: str = "2"
+    z: float = 0.0
+
+B1(1)
+B1("1")
+
+B1(1, "a")
+B1("a", 1)
+
+B1(1, "a", 1.0)
+B1("a", 1, "b")
+
+
+@dataclasses.dataclass(init=False)
+class B2:
+    x: int
+    y: str
+    z: float = 0.0
+
+    def __init__(self, x: int):
+        self.x = x
+        self.y = str(x)
+        self.z = 0.0
+
+B2(1)
+B2(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>)

python/testData/paramInfo/InitializingDataclass/a.py

@@ -6,4 +6,36 @@ class A:
     y: str
     z: float = 0.0
 
-A(<arg1>)
+A(<arg1>)
+
+
+@dataclasses.dataclass(init=True)
+class A2:
+    x: int
+    y: str
+    z: float = 0.0
+
+A2(<arg2>)
+
+
+@dataclasses.dataclass(init=False)
+class B1:
+    x: int = 1
+    y: str = "2"
+    z: float = 0.0
+
+B1(<arg3>)
+
+
+@dataclasses.dataclass(init=False)
+class B2:
+    x: int
+    y: str
+    z: float = 0.0
+
+    def __init__(self, x: int):
+        self.x = x
+        self.y = str(x)
+        self.z = 0.0
+
+B2(<arg4>)

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

@@ -712,9 +712,18 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
     runWithLanguageLevel(
       LanguageLevel.PYTHON37,
       () -> {
-        final Map<String, PsiElement> marks = loadMultiFileTest(1);
+        final Map<String, PsiElement> marks = loadMultiFileTest(4);
 
         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, "});
+
+        feignCtrlP(marks.get("<arg3>").getTextOffset()).check(
+          Arrays.asList("self: object", "cls: object"),
+          Arrays.asList(ArrayUtil.EMPTY_STRING_ARRAY, ArrayUtil.EMPTY_STRING_ARRAY),
+          Arrays.asList(new String[]{"self: object"}, new String[]{"cls: object"})
+        );
+
+        feignCtrlP(marks.get("<arg4>").getTextOffset()).check("self: B2, x: int", new String[]{"x: int"}, new String[]{"self: B2, "});
       }
     );
   }