diff --git a/ChangeLog b/ChangeLog index 6baadce493..bce6094b46 100644 --- a/ChangeLog +++ b/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. diff --git a/astroid/brain/brain_dataclasses.py b/astroid/brain/brain_dataclasses.py index bfdbbe09e5..a667e80df8 100644 --- a/astroid/brain/brain_dataclasses.py +++ b/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 diff --git a/astroid/nodes/scoped_nodes/scoped_nodes.py b/astroid/nodes/scoped_nodes/scoped_nodes.py index f8da9d2602..300f8c3371 100644 --- a/astroid/nodes/scoped_nodes/scoped_nodes.py +++ b/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, diff --git a/tests/unittest_brain_dataclasses.py b/tests/unittest_brain_dataclasses.py index 2054aebd62..6b33eec4a9 100644 --- a/tests/unittest_brain_dataclasses.py +++ b/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