From 610c1b70c30c5ac8cf3f448bb68b0c2e07b30902 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Thu, 15 Mar 2018 15:00:20 +0300
Subject: [PATCH] highlight `InitVar` and `ClassVar` with `default_factory`
 (PY-28957)

---
 .../inspections/PyDataclassInspection.kt      | 44 +++++++++++--------
 .../fieldDefaultAndDefaultFactory.py          |  9 +++-
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
index 2abaad0dddac..27472a3ea8e1 100644
--- a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
+++ b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
@@ -65,9 +65,12 @@ class PyDataclassInspection : PyInspection() {
           val initVars = mutableListOf<PyTargetExpression>()
 
           node.processClassLevelDeclarations { element, _ ->
-            if (element is PyTargetExpression && !PyTypingTypeProvider.isClassVar(element, myTypeEvalContext)) {
-              processDefaultFieldValue(element)
-              processAsInitVar(element, postInit)?.let { initVars.add(it) }
+            if (element is PyTargetExpression) {
+              if (!PyTypingTypeProvider.isClassVar(element, myTypeEvalContext)) {
+                processDefaultFieldValue(element)
+                processAsInitVar(element, postInit)?.let { initVars.add(it) }
+              }
+
               processFieldFunctionCall(element)
             }
 
@@ -149,16 +152,12 @@ class PyDataclassInspection : PyInspection() {
 
         if (parseDataclassParameters(cls, myTypeEvalContext) != null) {
           cls.processClassLevelDeclarations { element, _ ->
-            if (element is PyTargetExpression && element.name == node.name) {
-              val type = myTypeEvalContext.getType(element)
+            if (element is PyTargetExpression && element.name == node.name && isInitVar(element)) {
+              registerProblem(node.lastChild,
+                              "'${cls.name}' object could have no attribute '${element.name}' because it is declared as init-only",
+                              ProblemHighlightType.GENERIC_ERROR_OR_WARNING)
 
-              if (type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) {
-                registerProblem(node.lastChild,
-                                "'${cls.name}' object could have no attribute '${element.name}' because it is declared as init-only",
-                                ProblemHighlightType.GENERIC_ERROR_OR_WARNING)
-
-                return@processClassLevelDeclarations false
-              }
+              return@processClassLevelDeclarations false
             }
 
             true
@@ -258,9 +257,7 @@ class PyDataclassInspection : PyInspection() {
     }
 
     private fun processAsInitVar(field: PyTargetExpression, postInit: PyFunction?): PyTargetExpression? {
-      val type = myTypeEvalContext.getType(field)
-
-      if (type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) {
+      if (isInitVar(field)) {
         if (postInit == null) {
           registerProblem(field,
                           "Attribute '${field.name}' is useless until '$DUNDER_POST_INIT' is declared",
@@ -274,10 +271,17 @@ class PyDataclassInspection : PyInspection() {
     }
 
     private fun processFieldFunctionCall(field: PyTargetExpression) {
-      val fieldStub = PyDataclassFieldStubImpl.create(field)
-      if (fieldStub != null && fieldStub.hasDefault() && fieldStub.hasDefaultFactory()) {
-        val call = field.findAssignedValue() as? PyCallExpression ?: return
+      val fieldStub = PyDataclassFieldStubImpl.create(field) ?: return
+      val call = field.findAssignedValue() as? PyCallExpression ?: return
 
+      if (PyTypingTypeProvider.isClassVar(field, myTypeEvalContext) || isInitVar(field)) {
+        if (fieldStub.hasDefaultFactory()) {
+          registerProblem(call.getKeywordArgument("default_factory"),
+                          "Field cannot have a default factory",
+                          ProblemHighlightType.GENERIC_ERROR)
+        }
+      }
+      else if (fieldStub.hasDefault() && fieldStub.hasDefaultFactory()) {
         registerProblem(call.argumentList, "Cannot specify both 'default' and 'default_factory'", ProblemHighlightType.GENERIC_ERROR)
       }
     }
@@ -318,6 +322,10 @@ class PyDataclassInspection : PyInspection() {
       }
     }
 
+    private fun isInitVar(field: PyTargetExpression): Boolean {
+      return (myTypeEvalContext.getType(field) as? PyClassType)?.classQName == DATACLASSES_INITVAR_TYPE
+    }
+
     private fun isNotDataclass(type: PyType?, allowDefinition: Boolean): Boolean {
       if (type is PyStructuralType || PyTypeChecker.isUnknown(type, myTypeEvalContext)) return false
       if (type is PyUnionType) return type.members.all { isNotDataclass(it, allowDefinition) }
diff --git a/python/testData/inspections/PyDataclassInspection/fieldDefaultAndDefaultFactory.py b/python/testData/inspections/PyDataclassInspection/fieldDefaultAndDefaultFactory.py
index 0f936cd5ecbc..f0825e349a2d 100644
--- a/python/testData/inspections/PyDataclassInspection/fieldDefaultAndDefaultFactory.py
+++ b/python/testData/inspections/PyDataclassInspection/fieldDefaultAndDefaultFactory.py
@@ -1,7 +1,14 @@
 import dataclasses
+from typing import ClassVar
+
 
 @dataclasses.dataclass
 class E1:
     a: int = dataclasses.field(default=1)
     b: int = dataclasses.field(default_factory=int)
-    c: int = dataclasses.field<error descr="Cannot specify both 'default' and 'default_factory'">(default=1, default_factory=int)</error>
\ No newline at end of file
+    c: int = dataclasses.field<error descr="Cannot specify both 'default' and 'default_factory'">(default=1, default_factory=int)</error>
+    d: ClassVar[int] = dataclasses.field(default_factory=<error descr="Field cannot have a default factory">int</error>)
+    e: dataclasses.InitVar[int] = dataclasses.field(default_factory=<error descr="Field cannot have a default factory">int</error>)
+
+    def __post_init__(self, e: int):
+        pass
\ No newline at end of file