summaryrefslogtreecommitdiff
path: root/astroid
diff options
context:
space:
mode:
Diffstat (limited to 'astroid')
-rw-r--r--astroid/brain/brain_dataclasses.py50
-rw-r--r--astroid/tests/unittest_brain.py33
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()