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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
"""Utility functions for test code that uses astroid ASTs as input."""
import textwrap
from astroid import nodes
from astroid import builder
# The name of the transient function that is used to
# wrap expressions to be extracted when calling
# extract_node.
_TRANSIENT_FUNCTION = '__'
# The comment used to select a statement to be extracted
# when calling extract_node.
_STATEMENT_SELECTOR = '#@'
def _extract_expressions(node):
"""Find expressions in a call to _TRANSIENT_FUNCTION and extract them.
The function walks the AST recursively to search for expressions that
are wrapped into a call to _TRANSIENT_FUNCTION. If it finds such an
expression, it completely removes the function call node from the tree,
replacing it by the wrapped expression inside the parent.
:param node: An astroid node.
:type node: astroid.bases.NodeNG
:yields: The sequence of wrapped expressions on the modified tree
expression can be found.
"""
if (isinstance(node, nodes.CallFunc)
and isinstance(node.func, nodes.Name)
and node.func.name == _TRANSIENT_FUNCTION):
real_expr = node.args[0]
real_expr.parent = node.parent
# Search for node in all _astng_fields (the fields checked when
# get_children is called) of its parent. Some of those fields may
# be lists or tuples, in which case the elements need to be checked.
# When we find it, replace it by real_expr, so that the AST looks
# like no call to _TRANSIENT_FUNCTION ever took place.
for name in node.parent._astroid_fields:
child = getattr(node.parent, name)
if isinstance(child, (list, tuple)):
for idx, compound_child in enumerate(child):
if compound_child is node:
child[idx] = real_expr
elif child is node:
setattr(node.parent, name, real_expr)
yield real_expr
else:
for child in node.get_children():
for result in _extract_expressions(child):
yield result
def _find_statement_by_line(node, line):
"""Extracts the statement on a specific line from an AST.
If the line number of node matches line, it will be returned;
otherwise its children are iterated and the function is called
recursively.
:param node: An astroid node.
:type node: astroid.bases.NodeNG
:param line: The line number of the statement to extract.
:type line: int
:returns: The statement on the line, or None if no statement for the line
can be found.
:rtype: astroid.bases.NodeNG or None
"""
if isinstance(node, (nodes.Class, nodes.Function)):
# This is an inaccuracy in the AST: the nodes that can be
# decorated do not carry explicit information on which line
# the actual definition (class/def), but .fromline seems to
# be close enough.
node_line = node.fromlineno
else:
node_line = node.lineno
if node_line == line:
return node
for child in node.get_children():
result = _find_statement_by_line(child, line)
if result:
return result
return None
def extract_node(code, module_name=''):
"""Parses some Python code as a module and extracts a designated AST node.
Statements:
To extract one or more statement nodes, append #@ to the end of the line
Examples:
>>> def x():
>>> def y():
>>> return 1 #@
The return statement will be extracted.
>>> class X(object):
>>> def meth(self): #@
>>> pass
The funcion object 'meth' will be extracted.
Expressions:
To extract arbitrary expressions, surround them with the fake
function call __(...). After parsing, the surrounded expression
will be returned and the whole AST (accessible via the returned
node's parent attribute) will look like the function call was
never there in the first place.
Examples:
>>> a = __(1)
The const node will be extracted.
>>> def x(d=__(foo.bar)): pass
The node containing the default argument will be extracted.
>>> def foo(a, b):
>>> return 0 < __(len(a)) < b
The node containing the function call 'len' will be extracted.
If no statements or expressions are selected, the last toplevel
statement will be returned.
If the selected statement is a discard statement, (i.e. an expression
turned into a statement), the wrapped expression is returned instead.
For convenience, singleton lists are unpacked.
:param str code: A piece of Python code that is parsed as
a module. Will be passed through textwrap.dedent first.
:param str module_name: The name of the module.
:returns: The designated node from the parse tree, or a list of nodes.
:rtype: astroid.bases.NodeNG, or a list of nodes.
"""
def _extract(node):
if isinstance(node, nodes.Discard):
return node.value
else:
return node
requested_lines = []
for idx, line in enumerate(code.splitlines()):
if line.strip().endswith(_STATEMENT_SELECTOR):
requested_lines.append(idx + 1)
tree = build_module(code, module_name=module_name)
extracted = []
if requested_lines:
for line in requested_lines:
extracted.append(_find_statement_by_line(tree, line))
else:
# Modifies the tree.
extracted = list(_extract_expressions(tree))
if not extracted:
extracted.append(tree.body[-1])
extracted = [_extract(node) for node in extracted]
if len(extracted) == 1:
return extracted[0]
else:
return extracted
def build_module(code, module_name=''):
"""Parses a string module with a builder.
:param code: The code for the module.
:type code: str
:param module_name: The name for the module
:type module_name: str
:returns: The module AST.
:rtype: astroid.bases.NodeNG
"""
code = textwrap.dedent(code)
return builder.AstroidBuilder(None).string_build(code, modname=module_name)
|
