diff options
Diffstat (limited to 'astroid')
| -rw-r--r-- | astroid/brain/brain_dataclasses.py | 50 | ||||
| -rw-r--r-- | astroid/tests/unittest_brain.py | 33 |
2 files changed, 83 insertions, 0 deletions
diff --git a/astroid/brain/brain_dataclasses.py b/astroid/brain/brain_dataclasses.py new file mode 100644 index 00000000..7a25e0c6 --- /dev/null +++ b/astroid/brain/brain_dataclasses.py @@ -0,0 +1,50 @@ +# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html +# For details: https://github.com/PyCQA/astroid/blob/master/COPYING.LESSER +""" +Astroid hook for the dataclasses library +""" + +import astroid +from astroid import MANAGER + + +DATACLASSES_DECORATORS = frozenset(("dataclasses.dataclass", "dataclass")) + + +def is_decorated_with_dataclass(node, decorator_names=DATACLASSES_DECORATORS): + """Return True if a decorated node has a `dataclass` decorator applied.""" + if not node.decorators: + return False + for decorator_attribute in node.decorators.nodes: + if isinstance(decorator_attribute, astroid.Call): # decorator with arguments + decorator_attribute = decorator_attribute.func + if decorator_attribute.as_string() in decorator_names: + return True + return False + + +def dataclass_transform(node): + """Rewrite a dataclass to be easily understood by pylint""" + + for assign_node in node.body: + if not isinstance(assign_node, (astroid.AnnAssign, astroid.Assign)): + continue + + targets = ( + assign_node.targets + if hasattr(assign_node, "targets") + else [assign_node.target] + ) + for target in targets: + rhs_node = astroid.Unknown( + lineno=assign_node.lineno, + col_offset=assign_node.col_offset, + parent=assign_node, + ) + node.instance_attrs[target.name] = [rhs_node] + node.locals[target.name] = [rhs_node] + + +MANAGER.register_transform( + astroid.ClassDef, dataclass_transform, is_decorated_with_dataclass +) diff --git a/astroid/tests/unittest_brain.py b/astroid/tests/unittest_brain.py index a045b194..9b29a280 100644 --- a/astroid/tests/unittest_brain.py +++ b/astroid/tests/unittest_brain.py @@ -1927,5 +1927,38 @@ def test_crypt_brain(): assert attr in module +@pytest.mark.skipif(sys.version_info < (3, 7), reason="Dataclasses were added in 3.7") +def test_dataclasses(): + code = """ + import dataclasses + from dataclasses import dataclass + + @dataclass + class InventoryItem: + name: str + quantity_on_hand: int = 0 + + @dataclasses.dataclass + class Other: + name: str + """ + + module = astroid.parse(code) + first = module["InventoryItem"] + second = module["Other"] + + name = first.getattr("name") + assert len(name) == 1 + assert isinstance(name[0], astroid.Unknown) + + quantity_on_hand = first.getattr("quantity_on_hand") + assert len(quantity_on_hand) == 1 + assert isinstance(quantity_on_hand[0], astroid.Unknown) + + name = second.getattr("name") + assert len(name) == 1 + assert isinstance(name[0], astroid.Unknown) + + if __name__ == "__main__": unittest.main() |
