Skip to content

Add is_dataclass attribute to ClassDef #1391

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Feb 13, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
1 change: 1 addition & 0 deletions astroid/brain/brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion astroid/nodes/scoped_nodes/scoped_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down
31 changes: 31 additions & 0 deletions tests/unittest_brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You could also use extract_node together with #@. That would only return selected nodes.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll change it. I always use parse, but don't have a preference.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both work fine. extract_node might just be a bit easier to annotate if you only selected one node type.

"""
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