Check that object passed to attrs helpers is attrs (PY-26354)

python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt

@@ -22,6 +22,8 @@ class PyDataclassInspection : PyInspection() {
   companion object {
     private val ORDER_OPERATORS = setOf("__lt__", "__le__", "__gt__", "__ge__")
     private val DATACLASSES_HELPERS = setOf("dataclasses.fields", "dataclasses.asdict", "dataclasses.astuple", "dataclasses.replace")
+    private val ATTRS_HELPERS =
+      setOf("attr.__init__.fields", "attr.__init__.asdict", "attr.__init__.astuple", "attr.__init__.assoc", "attr.__init__.evolve")
   }
 
   override fun buildVisitor(holder: ProblemsHolder,
@@ -131,13 +133,24 @@ class PyDataclassInspection : PyInspection() {
         val callee = markedCallee?.element
         val calleeQName = callee?.qualifiedName
 
-        if (markedCallee != null && callee != null && DATACLASSES_HELPERS.contains(calleeQName)) {
+        if (markedCallee != null && callee != null) {
+          val dataclassType = when {
+            DATACLASSES_HELPERS.contains(calleeQName) -> PyDataclassParameters.Type.STD
+            ATTRS_HELPERS.contains(calleeQName) -> PyDataclassParameters.Type.ATTRS
+            else -> return
+          }
+
           val mapping = PyCallExpressionHelper.mapArguments(node, markedCallee, myTypeEvalContext)
 
           val dataclassParameter = callee.getParameters(myTypeEvalContext).firstOrNull()
           val dataclassArgument = mapping.mappedParameters.entries.firstOrNull { it.value == dataclassParameter }?.key
 
-          processHelperDataclassArgument(dataclassArgument, calleeQName!!)
+          if (dataclassType == PyDataclassParameters.Type.STD) {
+            processHelperDataclassArgument(dataclassArgument, calleeQName!!)
+          }
+          else if (dataclassType == PyDataclassParameters.Type.ATTRS) {
+            processHelperAttrsArgument(dataclassArgument, calleeQName!!)
+          }
         }
       }
     }
@@ -306,24 +319,41 @@ class PyDataclassInspection : PyInspection() {
 
       val allowDefinition = calleeQName == "dataclasses.fields"
 
-      if (isNotDataclass(myTypeEvalContext.getType(argument), allowDefinition)) {
+      if (isNotExpectedDataclass(myTypeEvalContext.getType(argument), PyDataclassParameters.Type.STD, allowDefinition, true)) {
         val message = "'$calleeQName' method should be called on dataclass instances" + if (allowDefinition) " or types" else ""
 
         registerProblem(argument, message)
       }
     }
 
+    private fun processHelperAttrsArgument(argument: PyExpression?, calleeQName: String) {
+      if (argument == null) return
+
+      val instance = calleeQName != "attr.__init__.fields"
+
+      if (isNotExpectedDataclass(myTypeEvalContext.getType(argument), PyDataclassParameters.Type.ATTRS, !instance, instance)) {
+        val presentableCalleeQName = calleeQName.replaceFirst(".__init__.", ".")
+        val message = "'$presentableCalleeQName' method should be called on attrs " + if (instance) "instances" else "types"
+
+        registerProblem(argument, message)
+      }
+    }
+
     private fun isInitVar(field: PyTargetExpression): Boolean {
       return (myTypeEvalContext.getType(field) as? PyClassType)?.classQName == DATACLASSES_INITVAR_TYPE
     }
 
-    private fun isNotDataclass(type: PyType?, allowDefinition: Boolean): Boolean {
+    private fun isNotExpectedDataclass(type: PyType?,
+                                       dataclassType: PyDataclassParameters.Type,
+                                       allowDefinition: Boolean,
+                                       allowInstance: Boolean): Boolean {
       if (type is PyStructuralType || PyTypeChecker.isUnknown(type, myTypeEvalContext)) return false
-      if (type is PyUnionType) return type.members.all { isNotDataclass(it, allowDefinition) }
+      if (type is PyUnionType) return type.members.all { isNotExpectedDataclass(it, dataclassType, allowDefinition, allowInstance) }
 
       return type !is PyClassType ||
              !allowDefinition && type.isDefinition ||
-             parseStdDataclassParameters(type.pyClass, myTypeEvalContext) == null
+             !allowInstance && !type.isDefinition ||
+             parseDataclassParameters(type.pyClass, myTypeEvalContext)?.type != dataclassType
     }
   }
 }

python/testData/inspections/PyDataclassInspection/attrsHelpersArgument.py

@@ -0,0 +1,69 @@
+import attr
+from typing import Type, Union
+
+
+class A:
+    pass
+
+
+attr.fields(<warning descr="'attr.fields' method should be called on attrs types">A</warning>)
+attr.fields(<warning descr="'attr.fields' method should be called on attrs types">A()</warning>)
+
+attr.asdict(<warning descr="'attr.asdict' method should be called on attrs instances">A()</warning>)
+attr.astuple(<warning descr="'attr.astuple' method should be called on attrs instances">A()</warning>)
+attr.assoc(<warning descr="'attr.assoc' method should be called on attrs instances">A()</warning>)
+attr.evolve(<warning descr="'attr.evolve' method should be called on attrs instances">A()</warning>)
+
+
+@attr.s
+class B:
+    pass
+
+
+attr.fields(B)
+attr.fields(<warning descr="'attr.fields' method should be called on attrs types">B()</warning>)
+
+attr.asdict(B())
+attr.astuple(B())
+attr.assoc(B())
+attr.evolve(B())
+
+attr.asdict(<warning descr="'attr.asdict' method should be called on attrs instances">B</warning>)
+attr.astuple(<warning descr="'attr.astuple' method should be called on attrs instances">B</warning>)
+attr.assoc(<warning descr="'attr.assoc' method should be called on attrs instances">B</warning>)
+attr.evolve(<warning descr="'attr.evolve' method should be called on attrs instances">B</warning>)
+
+
+def unknown(p):
+    attr.fields(p)
+
+    attr.asdict(p)
+    attr.astuple(p)
+
+
+def structural(p):
+    print(len(p))
+    attr.fields(p)
+
+    attr.asdict(p)
+    attr.astuple(p)
+    attr.assoc(p)
+    attr.evolve(p)
+
+
+def union1(p: Union[A, B]):
+    attr.fields(<warning descr="'attr.fields' method should be called on attrs types">p</warning>)
+
+    attr.asdict(p)
+    attr.astuple(p)
+    attr.assoc(p)
+    attr.evolve(p)
+
+
+def union2(p: Union[Type[A], Type[B]]):
+    attr.fields(p)
+
+    attr.asdict(<warning descr="'attr.asdict' method should be called on attrs instances">p</warning>)
+    attr.astuple(<warning descr="'attr.astuple' method should be called on attrs instances">p</warning>)
+    attr.assoc(<warning descr="'attr.assoc' method should be called on attrs instances">p</warning>)
+    attr.evolve(<warning descr="'attr.evolve' method should be called on attrs instances">p</warning>)

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

@@ -76,6 +76,11 @@ public class PyDataclassInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-26354
+  public void testAttrsHelpersArgument() {
+    doTest();
+  }
+
   // PY-27398
   public void testAccessToInitVar() {
     doTest();