From 0ea2055c4cc707618296b1294131a209bfdc4b0b Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Wed, 10 Jan 2018 15:22:16 +0300
Subject: [PATCH] Highlight wrong `__post_init__` signature (PY-27398)

---
 .../inspections/PyDataclassInspection.kt      | 38 +++++++++++++++++--
 .../wrongDunderPostInitSignature.py           | 19 ++++++++++
 .../PyDataclassInspectionTest.java            |  5 +++
 3 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 python/testData/inspections/PyDataclassInspection/wrongDunderPostInitSignature.py

diff --git a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
index fd0847f6c069..4f56bb9cda78 100644
--- a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
+++ b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
@@ -7,6 +7,7 @@ import com.intellij.codeInspection.LocalInspectionToolSession
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.psi.PsiElementVisitor
+import com.intellij.util.containers.ContainerUtil
 import com.jetbrains.python.codeInsight.stdlib.DATACLASSES_INITVAR_TYPE
 import com.jetbrains.python.codeInsight.stdlib.DUNDER_POST_INIT
 import com.jetbrains.python.codeInsight.stdlib.parseDataclassParameters
@@ -61,6 +62,8 @@ class PyDataclassInspection : PyInspection() {
                             ProblemHighlightType.LIKE_UNUSED_SYMBOL)
           }
 
+          val initVars = mutableListOf<String?>()
+
           node.processClassLevelDeclarations { element, _ ->
             if (element is PyTargetExpression && element.annotationValue != null) {
               val annotation = element.annotation
@@ -80,12 +83,17 @@ class PyDataclassInspection : PyInspection() {
                                     ProblemHighlightType.GENERIC_ERROR)
                   }
                 }
-                else if (postInit == null) {
+                else {
                   val type = myTypeEvalContext.getType(element)
                   if (type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) {
-                    registerProblem(element,
-                                    "attribute '${element.name}' is useless until '${DUNDER_POST_INIT}' is declared",
-                                    ProblemHighlightType.LIKE_UNUSED_SYMBOL)
+                    if (postInit == null) {
+                      registerProblem(element,
+                                      "attribute '${element.name}' is useless until '${DUNDER_POST_INIT}' is declared",
+                                      ProblemHighlightType.LIKE_UNUSED_SYMBOL)
+                    }
+                    else {
+                      initVars.add(element.name)
+                    }
                   }
                 }
               }
@@ -94,6 +102,28 @@ class PyDataclassInspection : PyInspection() {
             true
           }
 
+          if (postInit != null) {
+            val parameters = ContainerUtil.subList(postInit.getParameters(myTypeEvalContext), 1)
+
+            if (parameters.size != initVars.size) {
+              registerProblem(postInit.parameterList,
+                              "'${DUNDER_POST_INIT}' should take all init-only variables in the same order as they are defined",
+                              ProblemHighlightType.GENERIC_ERROR)
+            }
+            else {
+              parameters
+                .asSequence()
+                .zip(initVars.asSequence())
+                .all { it.first.name == it.second }
+                .also {
+                  if (!it) {
+                    registerProblem(postInit.parameterList,
+                                    "'${DUNDER_POST_INIT}' should take all init-only variables in the same order as they are defined")
+                  }
+                }
+            }
+          }
+
           PyNamedTupleInspection.inspectFieldsOrder(node, myTypeEvalContext, this::registerProblem)
         }
       }
diff --git a/python/testData/inspections/PyDataclassInspection/wrongDunderPostInitSignature.py b/python/testData/inspections/PyDataclassInspection/wrongDunderPostInitSignature.py
new file mode 100644
index 000000000000..373fff9b2b36
--- /dev/null
+++ b/python/testData/inspections/PyDataclassInspection/wrongDunderPostInitSignature.py
@@ -0,0 +1,19 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A:
+    a: int
+    b: dataclasses.InitVar[str]
+    c: dataclasses.InitVar[bytes]
+
+    def __post_init__<error descr="'__post_init__' should take all init-only variables in the same order as they are defined">(self)</error>:
+        pass
+
+@dataclasses.dataclass
+class B:
+    a: int
+    b: dataclasses.InitVar[str]
+    c: dataclasses.InitVar[bytes]
+
+    def __post_init__<warning descr="'__post_init__' should take all init-only variables in the same order as they are defined">(self, c, b)</warning>:
+        pass
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java
index 852ab256d193..acb0fb96a129 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java
@@ -66,6 +66,11 @@ public class PyDataclassInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-27398
+  public void testWrongDunderPostInitSignature() {
+    doTest();
+  }
+
   @Override
   protected void doTest() {
     runWithLanguageLevel(