forked from pixelneo/vim-python-docstring
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathasthelper.py
More file actions
130 lines (103 loc) · 3.82 KB
/
Copy pathasthelper.py
File metadata and controls
130 lines (103 loc) · 3.82 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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import ast
import sys
from itertools import chain
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 = {}
super().__init__()
def visit_Attribute(self, node):
if isinstance(node.value, ast.Name):
if node.value.id == self.instance_name:
self.data[node.attr] = None
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 = {}
self.instance_name = instance_name
def visit_Assign(self, node):
ac = AttributeCollector(self.instance_name)
for target in node.targets:
ac.visit(target)
self.attributes |= ac.data
def visit_AnnAssign(self, node):
ac = AttributeCollector(self.instance_name)
ac.visit(node.target)
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 chain(node.args.args, node.args.kwonlyargs):
type_hint = None
if arg.annotation is not None:
# ast.unparse doesn't work for python <= 3.8
if sys.version_info[0] == 3 and sys.version_info[1] <= 8:
from unparse import Unparser
from io import StringIO
v = StringIO()
Unparser(arg.annotation, file=v)
type_hint = v.getvalue()
else:
type_hint = ast.unparse(arg.annotation)
self.arguments.append({"arg": arg.arg, "type": type_hint})
if len(self.arguments) > 0 and (
self.arguments[0]["arg"] == "self" or self.arguments[0]["arg"] == "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)