Add is_dataclass attribute to ClassDef (#1391)

Author: DanielNoord

ChangeLog

@@ -30,6 +30,8 @@ Release date: TBA
 
   Closes #1330
 
+* Add ``is_dataclass`` attribute to ``ClassDef`` nodes.
+
 * Use ``sysconfig`` instead of ``distutils`` to determine the location of
   python stdlib files and packages.
 

astroid/brain/brain_dataclasses.py

@@ -67,6 +67,7 @@ def is_decorated_with_dataclass(node, decorator_names=DATACLASSES_DECORATORS):
 
 def dataclass_transform(node: ClassDef) -> None:
     """Rewrite a dataclass to be easily understood by pylint"""
+    node.is_dataclass = True
 
     for assign_node in _get_dataclass_attributes(node):
         name = assign_node.target.name

astroid/nodes/scoped_nodes/scoped_nodes.py

@@ -2134,7 +2134,7 @@ def my_meth(self, arg):
             ":type: str"
         ),
     )
-    _other_fields = ("name", "doc")
+    _other_fields = ("name", "doc", "is_dataclass")
     _other_other_fields = ("locals", "_newstyle")
     _newstyle = None
 
@@ -2212,6 +2212,9 @@ def __init__(
         :type doc: str or None
         """
 
+        self.is_dataclass: bool = False
+        """Whether this class is a dataclass."""
+
         super().__init__(
             lineno=lineno,
             col_offset=col_offset,

tests/unittest_brain_dataclasses.py

@@ -184,6 +184,7 @@ class A:
     inferred = next(class_def.infer())
     assert isinstance(inferred, nodes.ClassDef)
     assert inferred.instance_attrs == {}
+    assert inferred.is_dataclass
 
     # Both the class and instance can still access the attribute
     for node in (klass, instance):
@@ -216,6 +217,7 @@ class A:
     inferred = next(class_def.infer())
     assert isinstance(inferred, nodes.ClassDef)
     assert inferred.instance_attrs == {}
+    assert inferred.is_dataclass
 
     # Both the class and instance can still access the attribute
     for node in (klass, instance):
@@ -248,6 +250,7 @@ class A:
     inferred = next(class_def.infer())
     assert isinstance(inferred, nodes.ClassDef)
     assert inferred.instance_attrs == {}
+    assert inferred.is_dataclass
 
     # Both the class and instance can still access the attribute
     for node in (klass, instance):
@@ -666,6 +669,7 @@ class A:
     inferred = node.inferred()
     assert len(inferred) == 1 and isinstance(inferred[0], nodes.ClassDef)
     assert "attribute" in inferred[0].instance_attrs
+    assert inferred[0].is_dataclass
 
 
 @parametrize_module
@@ -683,3 +687,30 @@ class A:
     inferred = code.inferred()
     assert len(inferred) == 1
     assert isinstance(inferred[0], nodes.ClassDef)
+    assert inferred[0].is_dataclass
+
+
+def test_non_dataclass_is_not_dataclass() -> None:
+    """Test that something that isn't a dataclass has the correct attribute."""
+    module = astroid.parse(
+        """
+    class A:
+        val: field()
+
+    def dataclass():
+        return
+
+    @dataclass
+    class B:
+        val: field()
+    """
+    )
+    class_a = module.body[0].inferred()
+    assert len(class_a) == 1
+    assert isinstance(class_a[0], nodes.ClassDef)
+    assert not class_a[0].is_dataclass
+
+    class_b = module.body[2].inferred()
+    assert len(class_b) == 1
+    assert isinstance(class_b[0], nodes.ClassDef)
+    assert not class_b[0].is_dataclass