blob: 7a25e0c6361dc7cff277ae291b4fc29ac3fd6473 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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
)
|
