-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathasthelper.py
More file actions
112 lines (84 loc) · 2.99 KB
/
Copy pathasthelper.py
File metadata and controls
112 lines (84 loc) · 2.99 KB
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import ast
class RaiseNameCollector(ast.NodeVisitor):
def __init__(self):
self.data = set()
super().__init__()
def visit_Call(self, node):
self.data.add(node.func.id)
class AttributeCollector(ast.NodeVisitor):
def __init__(self, instance_name):
self.instance_name = instance_name
self.data = set()
super().__init__()
def visit_Attribute(self, node):
if isinstance(node.value, ast.Name):
if node.value.id == self.instance_name:
self.data.add(node.attr)
else:
self.generic_visit(node)
class ClassInstanceNameExtractor(ast.NodeVisitor):
def __init__(self):
self.instance_name = 'self' # default
self.set = False
super().__init__()
def visit_FunctionDef(self, node):
if node.name == '__init__':
self.instance_name = node.args.args[0].arg
self.set = True
elif not self.set:
self.instance_name = node.args.args[0].arg
def generic_visit(self, node):
if not self.set:
super().generic_visit(node)
class ClassVisitor(ast.NodeVisitor):
def __init__(self, instance_name):
super().__init__()
self.attributes = set()
self.instance_name = instance_name
def visit_Assign(self, node):
ac = AttributeCollector(self.instance_name)
for target in node.targets:
ac.visit(node)
self.attributes |= ac.data
class MethodVisitor(ast.NodeVisitor):
""" Gathers information about a method
Attributes:
arguments: arguments of the method
parent: indicated whether this method is inside another
raises: set of raised exceptions
returns: True if method returns
yields: True is method yields
"""
def __init__(self, parent=True):
self.parent = parent
self.arguments = []
self.raises = set()
self.returns = False
self.yields = False
super().__init__()
def _handle_functions(self, node):
new_visitor = MethodVisitor(parent=False)
new_visitor.generic_visit(node)
self.raises |= new_visitor.raises
if self.parent:
for arg in node.args.args:
self.arguments.append(arg.arg)
if len(self.arguments) > 0 and (self.arguments[0] == 'self' or self.arguments[0] == 'cls'):
self.arguments.pop(0)
self.returns = new_visitor.returns
self.yields = new_visitor.yields
def visit_Raise(self, node):
r = RaiseNameCollector()
r.visit(node)
self.raises |= r.data
super().generic_visit(node)
def visit_Yield(self, node):
self.yields = True
super().generic_visit(node)
def visit_Return(self, node):
self.returns = True
super().generic_visit(node)
def visit_FunctionDef(self, node):
self._handle_functions(node)
def visit_AsyncFunctionDef(self, node):
self._handle_functions(node)