From 574a48550808b76c4bdfc6ae0e5a17ba3a879982 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Fri, 4 May 2018 15:51:31 +0300
Subject: [PATCH] Highlight attrs attributes lacking type annotation (PY-26354)

---
 .../codeInsight/stdlib/PyDataclasses.kt       | 21 +++++++++++++------
 .../inspections/PyDataclassInspection.kt      | 13 ++++++++++++
 .../attrsLackingTypeAnnotation.py             | 15 +++++++++++++
 .../PyDataclassInspectionTest.java            |  5 +++++
 4 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 python/testData/inspections/PyDataclassInspection/attrsLackingTypeAnnotation.py

diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclasses.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclasses.kt
index 7c7932ed360f..0123f7a57b77 100644
--- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclasses.kt
+++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclasses.kt
@@ -3,11 +3,10 @@
  */
 package com.jetbrains.python.codeInsight.stdlib
 
-import com.jetbrains.python.psi.PyClass
-import com.jetbrains.python.psi.PyExpression
-import com.jetbrains.python.psi.PyKnownDecoratorUtil
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.PyKnownDecoratorUtil.KnownDecorator
-import com.jetbrains.python.psi.PyUtil
 import com.jetbrains.python.psi.impl.PyEvaluator
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.types.TypeEvalContext
@@ -47,7 +46,7 @@ private fun parseDataclassParameters(cls: PyClass,
       if (type != null) {
         for (mapping in decorator.multiMapArguments(PyResolveContext.noImplicits().withTypeEvalContext(context))) {
           if (mapping.unmappedArguments.isEmpty() && mapping.unmappedParameters.isEmpty()) {
-            val builder = PyDataclassParametersBuilder(type)
+            val builder = PyDataclassParametersBuilder(type, knownDecorator, cls)
 
             mapping
               .mappedParameters
@@ -87,7 +86,9 @@ data class PyDataclassParameters(val init: Boolean,
   }
 }
 
-private class PyDataclassParametersBuilder(private val type: PyDataclassParameters.Type) {
+private class PyDataclassParametersBuilder(private val type: PyDataclassParameters.Type,
+                                           decorator: KnownDecorator,
+                                           anchor: PsiElement) {
 
   companion object {
     private const val DEFAULT_INIT = true
@@ -114,6 +115,14 @@ private class PyDataclassParametersBuilder(private val type: PyDataclassParamete
 
   private val others = mutableMapOf<String, PyExpression>()
 
+  init {
+    if (type == PyDataclassParameters.Type.ATTRS && decorator == KnownDecorator.ATTR_DATACLASS) {
+      PyElementGenerator.getInstance(anchor.project)
+        .createExpressionFromText(LanguageLevel.forElement(anchor), PyNames.TRUE)
+        .also { others["auto_attribs"] = it }
+    }
+  }
+
   fun update(name: String?, argument: PyExpression?) {
     val value = PyUtil.peelArgument(argument)
 
diff --git a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
index fe7788e89839..645aad7e5030 100644
--- a/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
+++ b/python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt
@@ -93,6 +93,7 @@ class PyDataclassInspection : PyInspection() {
 
             processAttrsDefaultThroughDecorator(node)
             processAttrsInitializersAndValidators(node)
+            processAttrsAutoAttribs(node, dataclassParameters)
           }
 
           PyNamedTupleInspection.inspectFieldsOrder(
@@ -412,6 +413,18 @@ class PyDataclassInspection : PyInspection() {
       )
     }
 
+    private fun processAttrsAutoAttribs(cls: PyClass, dataclassParameters: PyDataclassParameters) {
+      if (PyEvaluator.evaluateAsBoolean(PyUtil.peelArgument(dataclassParameters.others["auto_attribs"]), false)) {
+        cls.processClassLevelDeclarations { element, _ ->
+          if (element is PyTargetExpression && element.annotation == null && PyDataclassFieldStubImpl.create(element) != null) {
+            registerProblem(element, "Attribute '${element.name}' lacks a type annotation", ProblemHighlightType.GENERIC_ERROR)
+          }
+
+          true
+        }
+      }
+    }
+
     private fun processAsInitVar(field: PyTargetExpression, postInit: PyFunction?): PyTargetExpression? {
       if (isInitVar(field)) {
         if (postInit == null) {
diff --git a/python/testData/inspections/PyDataclassInspection/attrsLackingTypeAnnotation.py b/python/testData/inspections/PyDataclassInspection/attrsLackingTypeAnnotation.py
new file mode 100644
index 000000000000..7cb3b5f84232
--- /dev/null
+++ b/python/testData/inspections/PyDataclassInspection/attrsLackingTypeAnnotation.py
@@ -0,0 +1,15 @@
+import attr
+
+@attr.dataclass
+class A1:
+    <error descr="Attribute 'a' lacks a type annotation">a</error> = attr.ib()
+    <error descr="Attribute 'b' lacks a type annotation">b</error> = attr.ib(type=int)
+    c = 1
+    d: int  = 1
+
+@attr.s(auto_attribs=True)
+class A3:
+    <error descr="Attribute 'a' lacks a type annotation">a</error> = attr.ib()
+    <error descr="Attribute 'b' lacks a type annotation">b</error> = attr.ib(type=int)
+    c = 1
+    d: int  = 1
\ 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 fa97b88cd5e2..96c34115f1ce 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/PyDataclassInspectionTest.java
@@ -166,6 +166,11 @@ public class PyDataclassInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-26354
+  public void testAttrsLackingTypeAnnotation() {
+    doTest();
+  }
+
   @Override
   protected void doTest() {
     runWithLanguageLevel(