blob: 02c3b2867fa85546d96c9a873d0428930b323e08 [file] [log] [blame]
Georg Brandl0c77a822008-06-10 16:37:50 +00001"""
2 ast
3 ~~~
4
5 The `ast` module helps Python applications to process trees of the Python
6 abstract syntax grammar. The abstract syntax itself might change with
7 each Python release; this module helps to find out programmatically what
8 the current grammar looks like and allows modifications of it.
9
10 An abstract syntax tree can be generated by passing `ast.PyCF_ONLY_AST` as
11 a flag to the `compile()` builtin function or by using the `parse()`
12 function from this module. The result will be a tree of objects whose
13 classes all inherit from `ast.AST`.
14
15 A modified abstract syntax tree can be compiled into a Python code object
16 using the built-in `compile()` function.
17
18 Additionally various helper functions are provided that make working with
19 the trees simpler. The main intention of the helper functions and this
20 module in general is to provide an easy to use interface for libraries
21 that work tightly with the python syntax (template engines for example).
22
23
24 :copyright: Copyright 2008 by Armin Ronacher.
25 :license: Python License.
26"""
27from _ast import *
28
29
Terry Reedyfeac6242011-01-24 21:36:03 +000030def parse(source, filename='<unknown>', mode='exec'):
Georg Brandl0c77a822008-06-10 16:37:50 +000031 """
Terry Reedyfeac6242011-01-24 21:36:03 +000032 Parse the source into an AST node.
33 Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
Georg Brandl0c77a822008-06-10 16:37:50 +000034 """
Terry Reedyfeac6242011-01-24 21:36:03 +000035 return compile(source, filename, mode, PyCF_ONLY_AST)
Georg Brandl0c77a822008-06-10 16:37:50 +000036
37
38def literal_eval(node_or_string):
39 """
40 Safely evaluate an expression node or a string containing a Python
41 expression. The string or node provided may only consist of the following
Éric Araujo2a83cc62011-04-17 19:10:27 +020042 Python literal structures: strings, bytes, numbers, tuples, lists, dicts,
43 sets, booleans, and None.
Georg Brandl0c77a822008-06-10 16:37:50 +000044 """
Georg Brandl0c77a822008-06-10 16:37:50 +000045 if isinstance(node_or_string, str):
46 node_or_string = parse(node_or_string, mode='eval')
47 if isinstance(node_or_string, Expression):
48 node_or_string = node_or_string.body
49 def _convert(node):
Benjamin Peterson5ef96e52010-07-11 23:06:06 +000050 if isinstance(node, (Str, Bytes)):
Georg Brandl0c77a822008-06-10 16:37:50 +000051 return node.s
52 elif isinstance(node, Num):
53 return node.n
54 elif isinstance(node, Tuple):
55 return tuple(map(_convert, node.elts))
56 elif isinstance(node, List):
57 return list(map(_convert, node.elts))
Georg Brandl492f3fc2010-07-11 09:41:21 +000058 elif isinstance(node, Set):
59 return set(map(_convert, node.elts))
Georg Brandl0c77a822008-06-10 16:37:50 +000060 elif isinstance(node, Dict):
61 return dict((_convert(k), _convert(v)) for k, v
62 in zip(node.keys, node.values))
Benjamin Peterson442f2092012-12-06 17:41:04 -050063 elif isinstance(node, NameConstant):
64 return node.value
Raymond Hettingerbc959732010-10-08 00:47:45 +000065 elif isinstance(node, UnaryOp) and \
66 isinstance(node.op, (UAdd, USub)) and \
67 isinstance(node.operand, (Num, UnaryOp, BinOp)):
68 operand = _convert(node.operand)
69 if isinstance(node.op, UAdd):
70 return + operand
71 else:
72 return - operand
Benjamin Peterson058e31e2009-01-16 03:54:08 +000073 elif isinstance(node, BinOp) and \
74 isinstance(node.op, (Add, Sub)) and \
Raymond Hettingerbc959732010-10-08 00:47:45 +000075 isinstance(node.right, (Num, UnaryOp, BinOp)) and \
76 isinstance(node.left, (Num, UnaryOp, BinOp)):
77 left = _convert(node.left)
78 right = _convert(node.right)
Benjamin Peterson058e31e2009-01-16 03:54:08 +000079 if isinstance(node.op, Add):
80 return left + right
81 else:
82 return left - right
Raymond Hettingerbc959732010-10-08 00:47:45 +000083 raise ValueError('malformed node or string: ' + repr(node))
Georg Brandl0c77a822008-06-10 16:37:50 +000084 return _convert(node_or_string)
85
86
87def dump(node, annotate_fields=True, include_attributes=False):
88 """
89 Return a formatted dump of the tree in *node*. This is mainly useful for
90 debugging purposes. The returned string will show the names and the values
91 for fields. This makes the code impossible to evaluate, so if evaluation is
92 wanted *annotate_fields* must be set to False. Attributes such as line
Benjamin Petersondcf97b92008-07-02 17:30:14 +000093 numbers and column offsets are not dumped by default. If this is wanted,
Georg Brandl0c77a822008-06-10 16:37:50 +000094 *include_attributes* can be set to True.
95 """
96 def _format(node):
97 if isinstance(node, AST):
98 fields = [(a, _format(b)) for a, b in iter_fields(node)]
99 rv = '%s(%s' % (node.__class__.__name__, ', '.join(
100 ('%s=%s' % field for field in fields)
101 if annotate_fields else
102 (b for a, b in fields)
103 ))
104 if include_attributes and node._attributes:
105 rv += fields and ', ' or ' '
106 rv += ', '.join('%s=%s' % (a, _format(getattr(node, a)))
107 for a in node._attributes)
108 return rv + ')'
109 elif isinstance(node, list):
110 return '[%s]' % ', '.join(_format(x) for x in node)
111 return repr(node)
112 if not isinstance(node, AST):
113 raise TypeError('expected AST, got %r' % node.__class__.__name__)
114 return _format(node)
115
116
117def copy_location(new_node, old_node):
118 """
119 Copy source location (`lineno` and `col_offset` attributes) from
120 *old_node* to *new_node* if possible, and return *new_node*.
121 """
122 for attr in 'lineno', 'col_offset':
123 if attr in old_node._attributes and attr in new_node._attributes \
124 and hasattr(old_node, attr):
125 setattr(new_node, attr, getattr(old_node, attr))
126 return new_node
127
128
129def fix_missing_locations(node):
130 """
131 When you compile a node tree with compile(), the compiler expects lineno and
132 col_offset attributes for every node that supports them. This is rather
133 tedious to fill in for generated nodes, so this helper adds these attributes
134 recursively where not already set, by setting them to the values of the
135 parent node. It works recursively starting at *node*.
136 """
137 def _fix(node, lineno, col_offset):
138 if 'lineno' in node._attributes:
139 if not hasattr(node, 'lineno'):
140 node.lineno = lineno
141 else:
142 lineno = node.lineno
143 if 'col_offset' in node._attributes:
144 if not hasattr(node, 'col_offset'):
145 node.col_offset = col_offset
146 else:
147 col_offset = node.col_offset
148 for child in iter_child_nodes(node):
149 _fix(child, lineno, col_offset)
150 _fix(node, 1, 0)
151 return node
152
153
154def increment_lineno(node, n=1):
155 """
156 Increment the line number of each node in the tree starting at *node* by *n*.
157 This is useful to "move code" to a different location in a file.
158 """
Georg Brandl0c77a822008-06-10 16:37:50 +0000159 for child in walk(node):
160 if 'lineno' in child._attributes:
161 child.lineno = getattr(child, 'lineno', 0) + n
162 return node
163
164
165def iter_fields(node):
166 """
167 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
168 that is present on *node*.
169 """
170 for field in node._fields:
171 try:
172 yield field, getattr(node, field)
173 except AttributeError:
174 pass
175
176
177def iter_child_nodes(node):
178 """
179 Yield all direct child nodes of *node*, that is, all fields that are nodes
180 and all items of fields that are lists of nodes.
181 """
182 for name, field in iter_fields(node):
183 if isinstance(field, AST):
184 yield field
185 elif isinstance(field, list):
186 for item in field:
187 if isinstance(item, AST):
188 yield item
189
190
191def get_docstring(node, clean=True):
192 """
193 Return the docstring for the given node or None if no docstring can
194 be found. If the node provided does not have docstrings a TypeError
195 will be raised.
196 """
197 if not isinstance(node, (FunctionDef, ClassDef, Module)):
198 raise TypeError("%r can't have docstrings" % node.__class__.__name__)
199 if node.body and isinstance(node.body[0], Expr) and \
200 isinstance(node.body[0].value, Str):
201 if clean:
202 import inspect
203 return inspect.cleandoc(node.body[0].value.s)
204 return node.body[0].value.s
205
206
207def walk(node):
208 """
Georg Brandl619e7ba2011-01-09 07:38:51 +0000209 Recursively yield all descendant nodes in the tree starting at *node*
210 (including *node* itself), in no specified order. This is useful if you
211 only want to modify nodes in place and don't care about the context.
Georg Brandl0c77a822008-06-10 16:37:50 +0000212 """
213 from collections import deque
214 todo = deque([node])
215 while todo:
216 node = todo.popleft()
217 todo.extend(iter_child_nodes(node))
218 yield node
219
220
221class NodeVisitor(object):
222 """
223 A node visitor base class that walks the abstract syntax tree and calls a
224 visitor function for every node found. This function may return a value
225 which is forwarded by the `visit` method.
226
227 This class is meant to be subclassed, with the subclass adding visitor
228 methods.
229
230 Per default the visitor functions for the nodes are ``'visit_'`` +
231 class name of the node. So a `TryFinally` node visit function would
232 be `visit_TryFinally`. This behavior can be changed by overriding
233 the `visit` method. If no visitor function exists for a node
234 (return value `None`) the `generic_visit` visitor is used instead.
235
236 Don't use the `NodeVisitor` if you want to apply changes to nodes during
237 traversing. For this a special visitor exists (`NodeTransformer`) that
238 allows modifications.
239 """
240
241 def visit(self, node):
242 """Visit a node."""
243 method = 'visit_' + node.__class__.__name__
244 visitor = getattr(self, method, self.generic_visit)
245 return visitor(node)
246
247 def generic_visit(self, node):
248 """Called if no explicit visitor function exists for a node."""
249 for field, value in iter_fields(node):
250 if isinstance(value, list):
251 for item in value:
252 if isinstance(item, AST):
253 self.visit(item)
254 elif isinstance(value, AST):
255 self.visit(value)
256
257
258class NodeTransformer(NodeVisitor):
259 """
260 A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
261 allows modification of nodes.
262
263 The `NodeTransformer` will walk the AST and use the return value of the
264 visitor methods to replace or remove the old node. If the return value of
265 the visitor method is ``None``, the node will be removed from its location,
266 otherwise it is replaced with the return value. The return value may be the
267 original node in which case no replacement takes place.
268
269 Here is an example transformer that rewrites all occurrences of name lookups
270 (``foo``) to ``data['foo']``::
271
272 class RewriteName(NodeTransformer):
273
274 def visit_Name(self, node):
275 return copy_location(Subscript(
276 value=Name(id='data', ctx=Load()),
277 slice=Index(value=Str(s=node.id)),
278 ctx=node.ctx
279 ), node)
280
281 Keep in mind that if the node you're operating on has child nodes you must
282 either transform the child nodes yourself or call the :meth:`generic_visit`
283 method for the node first.
284
285 For nodes that were part of a collection of statements (that applies to all
286 statement nodes), the visitor may also return a list of nodes rather than
287 just a single node.
288
289 Usually you use the transformer like this::
290
291 node = YourTransformer().visit(node)
292 """
293
294 def generic_visit(self, node):
295 for field, old_value in iter_fields(node):
296 old_value = getattr(node, field, None)
297 if isinstance(old_value, list):
298 new_values = []
299 for value in old_value:
300 if isinstance(value, AST):
301 value = self.visit(value)
302 if value is None:
303 continue
304 elif not isinstance(value, AST):
305 new_values.extend(value)
306 continue
307 new_values.append(value)
308 old_value[:] = new_values
309 elif isinstance(old_value, AST):
310 new_node = self.visit(old_value)
311 if new_node is None:
312 delattr(node, field)
313 else:
314 setattr(node, field, new_node)
315 return node