blob: 7a43581c0e6ce69cfde90cd6825fc43af058c045 [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"""
Pablo Galindo27fc3b62019-11-24 23:02:40 +000027import sys
Georg Brandl0c77a822008-06-10 16:37:50 +000028from _ast import *
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +030029from contextlib import contextmanager, nullcontext
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +030030from enum import IntEnum, auto
Georg Brandl0c77a822008-06-10 16:37:50 +000031
32
Guido van Rossum495da292019-03-07 12:38:08 -080033def parse(source, filename='<unknown>', mode='exec', *,
Guido van Rossum10b55c12019-06-11 17:23:12 -070034 type_comments=False, feature_version=None):
Georg Brandl0c77a822008-06-10 16:37:50 +000035 """
Terry Reedyfeac6242011-01-24 21:36:03 +000036 Parse the source into an AST node.
37 Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
Guido van Rossumdcfcd142019-01-31 03:40:27 -080038 Pass type_comments=True to get back type comments where the syntax allows.
Georg Brandl0c77a822008-06-10 16:37:50 +000039 """
Guido van Rossumdcfcd142019-01-31 03:40:27 -080040 flags = PyCF_ONLY_AST
41 if type_comments:
42 flags |= PyCF_TYPE_COMMENTS
Guido van Rossum10b55c12019-06-11 17:23:12 -070043 if isinstance(feature_version, tuple):
44 major, minor = feature_version # Should be a 2-tuple.
45 assert major == 3
46 feature_version = minor
47 elif feature_version is None:
48 feature_version = -1
49 # Else it should be an int giving the minor version for 3.x.
Guido van Rossum495da292019-03-07 12:38:08 -080050 return compile(source, filename, mode, flags,
Victor Stinnerefdf6ca2019-06-12 02:52:16 +020051 _feature_version=feature_version)
Georg Brandl0c77a822008-06-10 16:37:50 +000052
53
54def literal_eval(node_or_string):
55 """
56 Safely evaluate an expression node or a string containing a Python
57 expression. The string or node provided may only consist of the following
Éric Araujo2a83cc62011-04-17 19:10:27 +020058 Python literal structures: strings, bytes, numbers, tuples, lists, dicts,
59 sets, booleans, and None.
Georg Brandl0c77a822008-06-10 16:37:50 +000060 """
Georg Brandl0c77a822008-06-10 16:37:50 +000061 if isinstance(node_or_string, str):
62 node_or_string = parse(node_or_string, mode='eval')
63 if isinstance(node_or_string, Expression):
64 node_or_string = node_or_string.body
Curtis Bucherc21c5122020-05-05 12:40:56 -070065 def _raise_malformed_node(node):
66 raise ValueError(f'malformed node or string: {node!r}')
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020067 def _convert_num(node):
Curtis Bucherc21c5122020-05-05 12:40:56 -070068 if not isinstance(node, Constant) or type(node.value) not in (int, float, complex):
69 _raise_malformed_node(node)
70 return node.value
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020071 def _convert_signed_num(node):
72 if isinstance(node, UnaryOp) and isinstance(node.op, (UAdd, USub)):
73 operand = _convert_num(node.operand)
74 if isinstance(node.op, UAdd):
75 return + operand
76 else:
77 return - operand
78 return _convert_num(node)
Georg Brandl0c77a822008-06-10 16:37:50 +000079 def _convert(node):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010080 if isinstance(node, Constant):
81 return node.value
Georg Brandl0c77a822008-06-10 16:37:50 +000082 elif isinstance(node, Tuple):
83 return tuple(map(_convert, node.elts))
84 elif isinstance(node, List):
85 return list(map(_convert, node.elts))
Georg Brandl492f3fc2010-07-11 09:41:21 +000086 elif isinstance(node, Set):
87 return set(map(_convert, node.elts))
Raymond Hettinger4fcf5c12020-01-02 22:21:18 -070088 elif (isinstance(node, Call) and isinstance(node.func, Name) and
89 node.func.id == 'set' and node.args == node.keywords == []):
90 return set()
Georg Brandl0c77a822008-06-10 16:37:50 +000091 elif isinstance(node, Dict):
Curtis Bucherc21c5122020-05-05 12:40:56 -070092 if len(node.keys) != len(node.values):
93 _raise_malformed_node(node)
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020094 return dict(zip(map(_convert, node.keys),
95 map(_convert, node.values)))
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010096 elif isinstance(node, BinOp) and isinstance(node.op, (Add, Sub)):
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020097 left = _convert_signed_num(node.left)
98 right = _convert_num(node.right)
99 if isinstance(left, (int, float)) and isinstance(right, complex):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +0100100 if isinstance(node.op, Add):
101 return left + right
102 else:
103 return left - right
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +0200104 return _convert_signed_num(node)
Georg Brandl0c77a822008-06-10 16:37:50 +0000105 return _convert(node_or_string)
106
107
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300108def dump(node, annotate_fields=True, include_attributes=False, *, indent=None):
Georg Brandl0c77a822008-06-10 16:37:50 +0000109 """
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300110 Return a formatted dump of the tree in node. This is mainly useful for
111 debugging purposes. If annotate_fields is true (by default),
112 the returned string will show the names and the values for fields.
113 If annotate_fields is false, the result string will be more compact by
114 omitting unambiguous field names. Attributes such as line
Benjamin Petersondcf97b92008-07-02 17:30:14 +0000115 numbers and column offsets are not dumped by default. If this is wanted,
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300116 include_attributes can be set to true. If indent is a non-negative
117 integer or string, then the tree will be pretty-printed with that indent
118 level. None (the default) selects the single line representation.
Georg Brandl0c77a822008-06-10 16:37:50 +0000119 """
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300120 def _format(node, level=0):
121 if indent is not None:
122 level += 1
123 prefix = '\n' + indent * level
124 sep = ',\n' + indent * level
125 else:
126 prefix = ''
127 sep = ', '
Georg Brandl0c77a822008-06-10 16:37:50 +0000128 if isinstance(node, AST):
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200129 cls = type(node)
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300130 args = []
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300131 allsimple = True
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300132 keywords = annotate_fields
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200133 for name in node._fields:
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300134 try:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200135 value = getattr(node, name)
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300136 except AttributeError:
137 keywords = True
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200138 continue
139 if value is None and getattr(cls, name, ...) is None:
140 keywords = True
141 continue
142 value, simple = _format(value, level)
143 allsimple = allsimple and simple
144 if keywords:
145 args.append('%s=%s' % (name, value))
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300146 else:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200147 args.append(value)
148 if include_attributes and node._attributes:
149 for name in node._attributes:
150 try:
151 value = getattr(node, name)
152 except AttributeError:
153 continue
154 if value is None and getattr(cls, name, ...) is None:
155 continue
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300156 value, simple = _format(value, level)
157 allsimple = allsimple and simple
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200158 args.append('%s=%s' % (name, value))
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300159 if allsimple and len(args) <= 3:
160 return '%s(%s)' % (node.__class__.__name__, ', '.join(args)), not args
161 return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args)), False
Georg Brandl0c77a822008-06-10 16:37:50 +0000162 elif isinstance(node, list):
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300163 if not node:
164 return '[]', True
165 return '[%s%s]' % (prefix, sep.join(_format(x, level)[0] for x in node)), False
166 return repr(node), True
167
Georg Brandl0c77a822008-06-10 16:37:50 +0000168 if not isinstance(node, AST):
169 raise TypeError('expected AST, got %r' % node.__class__.__name__)
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300170 if indent is not None and not isinstance(indent, str):
171 indent = ' ' * indent
172 return _format(node)[0]
Georg Brandl0c77a822008-06-10 16:37:50 +0000173
174
175def copy_location(new_node, old_node):
176 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000177 Copy source location (`lineno`, `col_offset`, `end_lineno`, and `end_col_offset`
178 attributes) from *old_node* to *new_node* if possible, and return *new_node*.
Georg Brandl0c77a822008-06-10 16:37:50 +0000179 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000180 for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200181 if attr in old_node._attributes and attr in new_node._attributes:
182 value = getattr(old_node, attr, None)
183 if value is not None:
184 setattr(new_node, attr, value)
Georg Brandl0c77a822008-06-10 16:37:50 +0000185 return new_node
186
187
188def fix_missing_locations(node):
189 """
190 When you compile a node tree with compile(), the compiler expects lineno and
191 col_offset attributes for every node that supports them. This is rather
192 tedious to fill in for generated nodes, so this helper adds these attributes
193 recursively where not already set, by setting them to the values of the
194 parent node. It works recursively starting at *node*.
195 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000196 def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
Georg Brandl0c77a822008-06-10 16:37:50 +0000197 if 'lineno' in node._attributes:
198 if not hasattr(node, 'lineno'):
199 node.lineno = lineno
200 else:
201 lineno = node.lineno
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000202 if 'end_lineno' in node._attributes:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200203 if getattr(node, 'end_lineno', None) is None:
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000204 node.end_lineno = end_lineno
205 else:
206 end_lineno = node.end_lineno
Georg Brandl0c77a822008-06-10 16:37:50 +0000207 if 'col_offset' in node._attributes:
208 if not hasattr(node, 'col_offset'):
209 node.col_offset = col_offset
210 else:
211 col_offset = node.col_offset
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000212 if 'end_col_offset' in node._attributes:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200213 if getattr(node, 'end_col_offset', None) is None:
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000214 node.end_col_offset = end_col_offset
215 else:
216 end_col_offset = node.end_col_offset
Georg Brandl0c77a822008-06-10 16:37:50 +0000217 for child in iter_child_nodes(node):
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000218 _fix(child, lineno, col_offset, end_lineno, end_col_offset)
219 _fix(node, 1, 0, 1, 0)
Georg Brandl0c77a822008-06-10 16:37:50 +0000220 return node
221
222
223def increment_lineno(node, n=1):
224 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000225 Increment the line number and end line number of each node in the tree
226 starting at *node* by *n*. This is useful to "move code" to a different
227 location in a file.
Georg Brandl0c77a822008-06-10 16:37:50 +0000228 """
Georg Brandl0c77a822008-06-10 16:37:50 +0000229 for child in walk(node):
230 if 'lineno' in child._attributes:
231 child.lineno = getattr(child, 'lineno', 0) + n
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000232 if 'end_lineno' in child._attributes:
233 child.end_lineno = getattr(child, 'end_lineno', 0) + n
Georg Brandl0c77a822008-06-10 16:37:50 +0000234 return node
235
236
237def iter_fields(node):
238 """
239 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
240 that is present on *node*.
241 """
242 for field in node._fields:
243 try:
244 yield field, getattr(node, field)
245 except AttributeError:
246 pass
247
248
249def iter_child_nodes(node):
250 """
251 Yield all direct child nodes of *node*, that is, all fields that are nodes
252 and all items of fields that are lists of nodes.
253 """
254 for name, field in iter_fields(node):
255 if isinstance(field, AST):
256 yield field
257 elif isinstance(field, list):
258 for item in field:
259 if isinstance(item, AST):
260 yield item
261
262
263def get_docstring(node, clean=True):
264 """
265 Return the docstring for the given node or None if no docstring can
266 be found. If the node provided does not have docstrings a TypeError
267 will be raised.
Matthias Bussonnier41cea702017-02-23 22:44:19 -0800268
269 If *clean* is `True`, all tabs are expanded to spaces and any whitespace
270 that can be uniformly removed from the second line onwards is removed.
Georg Brandl0c77a822008-06-10 16:37:50 +0000271 """
Yury Selivanov2f07a662015-07-23 08:54:35 +0300272 if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
Georg Brandl0c77a822008-06-10 16:37:50 +0000273 raise TypeError("%r can't have docstrings" % node.__class__.__name__)
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300274 if not(node.body and isinstance(node.body[0], Expr)):
Serhiy Storchaka73cbe7a2018-05-29 12:04:55 +0300275 return None
276 node = node.body[0].value
277 if isinstance(node, Str):
278 text = node.s
279 elif isinstance(node, Constant) and isinstance(node.value, str):
280 text = node.value
281 else:
282 return None
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300283 if clean:
Victor Stinnerf2c1aa12016-01-26 00:40:57 +0100284 import inspect
285 text = inspect.cleandoc(text)
286 return text
Georg Brandl0c77a822008-06-10 16:37:50 +0000287
288
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000289def _splitlines_no_ff(source):
290 """Split a string into lines ignoring form feed and other chars.
291
292 This mimics how the Python parser splits source code.
293 """
294 idx = 0
295 lines = []
296 next_line = ''
297 while idx < len(source):
298 c = source[idx]
299 next_line += c
300 idx += 1
301 # Keep \r\n together
302 if c == '\r' and idx < len(source) and source[idx] == '\n':
303 next_line += '\n'
304 idx += 1
305 if c in '\r\n':
306 lines.append(next_line)
307 next_line = ''
308
309 if next_line:
310 lines.append(next_line)
311 return lines
312
313
314def _pad_whitespace(source):
mpheathfbeba8f2020-02-14 04:32:09 +1000315 r"""Replace all chars except '\f\t' in a line with spaces."""
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000316 result = ''
317 for c in source:
318 if c in '\f\t':
319 result += c
320 else:
321 result += ' '
322 return result
323
324
325def get_source_segment(source, node, *, padded=False):
326 """Get source code segment of the *source* that generated *node*.
327
328 If some location information (`lineno`, `end_lineno`, `col_offset`,
329 or `end_col_offset`) is missing, return None.
330
331 If *padded* is `True`, the first line of a multi-line statement will
332 be padded with spaces to match its original position.
333 """
334 try:
335 lineno = node.lineno - 1
336 end_lineno = node.end_lineno - 1
337 col_offset = node.col_offset
338 end_col_offset = node.end_col_offset
339 except AttributeError:
340 return None
341
342 lines = _splitlines_no_ff(source)
343 if end_lineno == lineno:
344 return lines[lineno].encode()[col_offset:end_col_offset].decode()
345
346 if padded:
347 padding = _pad_whitespace(lines[lineno].encode()[:col_offset].decode())
348 else:
349 padding = ''
350
351 first = padding + lines[lineno].encode()[col_offset:].decode()
352 last = lines[end_lineno].encode()[:end_col_offset].decode()
353 lines = lines[lineno+1:end_lineno]
354
355 lines.insert(0, first)
356 lines.append(last)
357 return ''.join(lines)
358
359
Georg Brandl0c77a822008-06-10 16:37:50 +0000360def walk(node):
361 """
Georg Brandl619e7ba2011-01-09 07:38:51 +0000362 Recursively yield all descendant nodes in the tree starting at *node*
363 (including *node* itself), in no specified order. This is useful if you
364 only want to modify nodes in place and don't care about the context.
Georg Brandl0c77a822008-06-10 16:37:50 +0000365 """
366 from collections import deque
367 todo = deque([node])
368 while todo:
369 node = todo.popleft()
370 todo.extend(iter_child_nodes(node))
371 yield node
372
373
374class NodeVisitor(object):
375 """
376 A node visitor base class that walks the abstract syntax tree and calls a
377 visitor function for every node found. This function may return a value
378 which is forwarded by the `visit` method.
379
380 This class is meant to be subclassed, with the subclass adding visitor
381 methods.
382
383 Per default the visitor functions for the nodes are ``'visit_'`` +
384 class name of the node. So a `TryFinally` node visit function would
385 be `visit_TryFinally`. This behavior can be changed by overriding
386 the `visit` method. If no visitor function exists for a node
387 (return value `None`) the `generic_visit` visitor is used instead.
388
389 Don't use the `NodeVisitor` if you want to apply changes to nodes during
390 traversing. For this a special visitor exists (`NodeTransformer`) that
391 allows modifications.
392 """
393
394 def visit(self, node):
395 """Visit a node."""
396 method = 'visit_' + node.__class__.__name__
397 visitor = getattr(self, method, self.generic_visit)
398 return visitor(node)
399
400 def generic_visit(self, node):
401 """Called if no explicit visitor function exists for a node."""
402 for field, value in iter_fields(node):
403 if isinstance(value, list):
404 for item in value:
405 if isinstance(item, AST):
406 self.visit(item)
407 elif isinstance(value, AST):
408 self.visit(value)
409
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300410 def visit_Constant(self, node):
411 value = node.value
412 type_name = _const_node_type_names.get(type(value))
413 if type_name is None:
414 for cls, name in _const_node_type_names.items():
415 if isinstance(value, cls):
416 type_name = name
417 break
418 if type_name is not None:
419 method = 'visit_' + type_name
420 try:
421 visitor = getattr(self, method)
422 except AttributeError:
423 pass
424 else:
425 import warnings
426 warnings.warn(f"{method} is deprecated; add visit_Constant",
427 DeprecationWarning, 2)
428 return visitor(node)
429 return self.generic_visit(node)
430
Georg Brandl0c77a822008-06-10 16:37:50 +0000431
432class NodeTransformer(NodeVisitor):
433 """
434 A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
435 allows modification of nodes.
436
437 The `NodeTransformer` will walk the AST and use the return value of the
438 visitor methods to replace or remove the old node. If the return value of
439 the visitor method is ``None``, the node will be removed from its location,
440 otherwise it is replaced with the return value. The return value may be the
441 original node in which case no replacement takes place.
442
443 Here is an example transformer that rewrites all occurrences of name lookups
444 (``foo``) to ``data['foo']``::
445
446 class RewriteName(NodeTransformer):
447
448 def visit_Name(self, node):
Pablo Galindoc00c86b2020-03-12 00:48:19 +0000449 return Subscript(
Georg Brandl0c77a822008-06-10 16:37:50 +0000450 value=Name(id='data', ctx=Load()),
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200451 slice=Constant(value=node.id),
Georg Brandl0c77a822008-06-10 16:37:50 +0000452 ctx=node.ctx
Pablo Galindoc00c86b2020-03-12 00:48:19 +0000453 )
Georg Brandl0c77a822008-06-10 16:37:50 +0000454
455 Keep in mind that if the node you're operating on has child nodes you must
456 either transform the child nodes yourself or call the :meth:`generic_visit`
457 method for the node first.
458
459 For nodes that were part of a collection of statements (that applies to all
460 statement nodes), the visitor may also return a list of nodes rather than
461 just a single node.
462
463 Usually you use the transformer like this::
464
465 node = YourTransformer().visit(node)
466 """
467
468 def generic_visit(self, node):
469 for field, old_value in iter_fields(node):
Georg Brandl0c77a822008-06-10 16:37:50 +0000470 if isinstance(old_value, list):
471 new_values = []
472 for value in old_value:
473 if isinstance(value, AST):
474 value = self.visit(value)
475 if value is None:
476 continue
477 elif not isinstance(value, AST):
478 new_values.extend(value)
479 continue
480 new_values.append(value)
481 old_value[:] = new_values
482 elif isinstance(old_value, AST):
483 new_node = self.visit(old_value)
484 if new_node is None:
485 delattr(node, field)
486 else:
487 setattr(node, field, new_node)
488 return node
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300489
490
491# The following code is for backward compatibility.
492# It will be removed in future.
493
494def _getter(self):
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200495 """Deprecated. Use value instead."""
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300496 return self.value
497
498def _setter(self, value):
499 self.value = value
500
501Constant.n = property(_getter, _setter)
502Constant.s = property(_getter, _setter)
503
504class _ABC(type):
505
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200506 def __init__(cls, *args):
507 cls.__doc__ = """Deprecated AST node class. Use ast.Constant instead"""
508
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300509 def __instancecheck__(cls, inst):
510 if not isinstance(inst, Constant):
511 return False
512 if cls in _const_types:
513 try:
514 value = inst.value
515 except AttributeError:
516 return False
517 else:
Anthony Sottile74176222019-01-18 11:30:28 -0800518 return (
519 isinstance(value, _const_types[cls]) and
520 not isinstance(value, _const_types_not.get(cls, ()))
521 )
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300522 return type.__instancecheck__(cls, inst)
523
524def _new(cls, *args, **kwargs):
525 if cls in _const_types:
526 return Constant(*args, **kwargs)
527 return Constant.__new__(cls, *args, **kwargs)
528
529class Num(Constant, metaclass=_ABC):
530 _fields = ('n',)
531 __new__ = _new
532
533class Str(Constant, metaclass=_ABC):
534 _fields = ('s',)
535 __new__ = _new
536
537class Bytes(Constant, metaclass=_ABC):
538 _fields = ('s',)
539 __new__ = _new
540
541class NameConstant(Constant, metaclass=_ABC):
542 __new__ = _new
543
544class Ellipsis(Constant, metaclass=_ABC):
545 _fields = ()
546
547 def __new__(cls, *args, **kwargs):
548 if cls is Ellipsis:
549 return Constant(..., *args, **kwargs)
550 return Constant.__new__(cls, *args, **kwargs)
551
552_const_types = {
553 Num: (int, float, complex),
554 Str: (str,),
555 Bytes: (bytes,),
556 NameConstant: (type(None), bool),
557 Ellipsis: (type(...),),
558}
Anthony Sottile74176222019-01-18 11:30:28 -0800559_const_types_not = {
560 Num: (bool,),
561}
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200562
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300563_const_node_type_names = {
564 bool: 'NameConstant', # should be before int
565 type(None): 'NameConstant',
566 int: 'Num',
567 float: 'Num',
568 complex: 'Num',
569 str: 'Str',
570 bytes: 'Bytes',
571 type(...): 'Ellipsis',
572}
Serhiy Storchaka832e8642019-09-09 23:36:13 +0300573
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200574class slice(AST):
575 """Deprecated AST node class."""
576
577class Index(slice):
578 """Deprecated AST node class. Use the index value directly instead."""
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200579 def __new__(cls, value, **kwargs):
580 return value
581
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200582class ExtSlice(slice):
583 """Deprecated AST node class. Use ast.Tuple instead."""
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200584 def __new__(cls, dims=(), **kwargs):
585 return Tuple(list(dims), Load(), **kwargs)
586
587def _dims_getter(self):
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200588 """Deprecated. Use elts instead."""
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200589 return self.elts
590
591def _dims_setter(self, value):
592 self.elts = value
593
594Tuple.dims = property(_dims_getter, _dims_setter)
595
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200596class Suite(mod):
597 """Deprecated AST node class. Unused in Python 3."""
598
599class AugLoad(expr_context):
600 """Deprecated AST node class. Unused in Python 3."""
601
602class AugStore(expr_context):
603 """Deprecated AST node class. Unused in Python 3."""
604
605class Param(expr_context):
606 """Deprecated AST node class. Unused in Python 3."""
607
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200608
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000609# Large float and imaginary literals get turned into infinities in the AST.
610# We unparse those infinities to INFSTR.
611_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
612
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300613class _Precedence(IntEnum):
614 """Precedence table that originated from python grammar."""
615
616 TUPLE = auto()
617 YIELD = auto() # 'yield', 'yield from'
618 TEST = auto() # 'if'-'else', 'lambda'
619 OR = auto() # 'or'
620 AND = auto() # 'and'
621 NOT = auto() # 'not'
622 CMP = auto() # '<', '>', '==', '>=', '<=', '!=',
623 # 'in', 'not in', 'is', 'is not'
624 EXPR = auto()
625 BOR = EXPR # '|'
626 BXOR = auto() # '^'
627 BAND = auto() # '&'
628 SHIFT = auto() # '<<', '>>'
629 ARITH = auto() # '+', '-'
630 TERM = auto() # '*', '@', '/', '%', '//'
631 FACTOR = auto() # unary '+', '-', '~'
632 POWER = auto() # '**'
633 AWAIT = auto() # 'await'
634 ATOM = auto()
635
636 def next(self):
637 try:
638 return self.__class__(self + 1)
639 except ValueError:
640 return self
641
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000642class _Unparser(NodeVisitor):
643 """Methods in this class recursively traverse an AST and
644 output source code for the abstract syntax; original formatting
645 is disregarded."""
646
647 def __init__(self):
648 self._source = []
649 self._buffer = []
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300650 self._precedences = {}
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000651 self._indent = 0
652
653 def interleave(self, inter, f, seq):
654 """Call f on each item in seq, calling inter() in between."""
655 seq = iter(seq)
656 try:
657 f(next(seq))
658 except StopIteration:
659 pass
660 else:
661 for x in seq:
662 inter()
663 f(x)
664
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +0300665 def items_view(self, traverser, items):
666 """Traverse and separate the given *items* with a comma and append it to
667 the buffer. If *items* is a single item sequence, a trailing comma
668 will be added."""
669 if len(items) == 1:
670 traverser(items[0])
671 self.write(",")
672 else:
673 self.interleave(lambda: self.write(", "), traverser, items)
674
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300675 def maybe_newline(self):
676 """Adds a newline if it isn't the start of generated source"""
677 if self._source:
678 self.write("\n")
679
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000680 def fill(self, text=""):
681 """Indent a piece of text and append it, according to the current
682 indentation level"""
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300683 self.maybe_newline()
684 self.write(" " * self._indent + text)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000685
686 def write(self, text):
687 """Append a piece of text"""
688 self._source.append(text)
689
690 def buffer_writer(self, text):
691 self._buffer.append(text)
692
693 @property
694 def buffer(self):
695 value = "".join(self._buffer)
696 self._buffer.clear()
697 return value
698
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000699 @contextmanager
700 def block(self):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000701 """A context manager for preparing the source for blocks. It adds
702 the character':', increases the indentation on enter and decreases
703 the indentation on exit."""
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000704 self.write(":")
705 self._indent += 1
706 yield
707 self._indent -= 1
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000708
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300709 @contextmanager
710 def delimit(self, start, end):
711 """A context manager for preparing the source for expressions. It adds
712 *start* to the buffer and enters, after exit it adds *end*."""
713
714 self.write(start)
715 yield
716 self.write(end)
717
718 def delimit_if(self, start, end, condition):
719 if condition:
720 return self.delimit(start, end)
721 else:
722 return nullcontext()
723
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300724 def require_parens(self, precedence, node):
725 """Shortcut to adding precedence related parens"""
726 return self.delimit_if("(", ")", self.get_precedence(node) > precedence)
727
728 def get_precedence(self, node):
729 return self._precedences.get(node, _Precedence.TEST)
730
731 def set_precedence(self, precedence, *nodes):
732 for node in nodes:
733 self._precedences[node] = precedence
734
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300735 def get_raw_docstring(self, node):
736 """If a docstring node is found in the body of the *node* parameter,
737 return that docstring node, None otherwise.
738
739 Logic mirrored from ``_PyAST_GetDocString``."""
740 if not isinstance(
741 node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
742 ) or len(node.body) < 1:
743 return None
744 node = node.body[0]
745 if not isinstance(node, Expr):
746 return None
747 node = node.value
748 if isinstance(node, Constant) and isinstance(node.value, str):
749 return node
750
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000751 def traverse(self, node):
752 if isinstance(node, list):
753 for item in node:
754 self.traverse(item)
755 else:
756 super().visit(node)
757
758 def visit(self, node):
759 """Outputs a source code string that, if converted back to an ast
760 (using ast.parse) will generate an AST equivalent to *node*"""
761 self._source = []
762 self.traverse(node)
763 return "".join(self._source)
764
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300765 def _write_docstring_and_traverse_body(self, node):
766 if (docstring := self.get_raw_docstring(node)):
767 self._write_docstring(docstring)
768 self.traverse(node.body[1:])
769 else:
770 self.traverse(node.body)
771
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000772 def visit_Module(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300773 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000774
Batuhan Taşkaya5b66ec12020-03-15 22:56:57 +0300775 def visit_FunctionType(self, node):
776 with self.delimit("(", ")"):
777 self.interleave(
778 lambda: self.write(", "), self.traverse, node.argtypes
779 )
780
781 self.write(" -> ")
782 self.traverse(node.returns)
783
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000784 def visit_Expr(self, node):
785 self.fill()
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300786 self.set_precedence(_Precedence.YIELD, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000787 self.traverse(node.value)
788
789 def visit_NamedExpr(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300790 with self.require_parens(_Precedence.TUPLE, node):
791 self.set_precedence(_Precedence.ATOM, node.target, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300792 self.traverse(node.target)
793 self.write(" := ")
794 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000795
796 def visit_Import(self, node):
797 self.fill("import ")
798 self.interleave(lambda: self.write(", "), self.traverse, node.names)
799
800 def visit_ImportFrom(self, node):
801 self.fill("from ")
802 self.write("." * node.level)
803 if node.module:
804 self.write(node.module)
805 self.write(" import ")
806 self.interleave(lambda: self.write(", "), self.traverse, node.names)
807
808 def visit_Assign(self, node):
809 self.fill()
810 for target in node.targets:
811 self.traverse(target)
812 self.write(" = ")
813 self.traverse(node.value)
814
815 def visit_AugAssign(self, node):
816 self.fill()
817 self.traverse(node.target)
818 self.write(" " + self.binop[node.op.__class__.__name__] + "= ")
819 self.traverse(node.value)
820
821 def visit_AnnAssign(self, node):
822 self.fill()
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300823 with self.delimit_if("(", ")", not node.simple and isinstance(node.target, Name)):
824 self.traverse(node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000825 self.write(": ")
826 self.traverse(node.annotation)
827 if node.value:
828 self.write(" = ")
829 self.traverse(node.value)
830
831 def visit_Return(self, node):
832 self.fill("return")
833 if node.value:
834 self.write(" ")
835 self.traverse(node.value)
836
837 def visit_Pass(self, node):
838 self.fill("pass")
839
840 def visit_Break(self, node):
841 self.fill("break")
842
843 def visit_Continue(self, node):
844 self.fill("continue")
845
846 def visit_Delete(self, node):
847 self.fill("del ")
848 self.interleave(lambda: self.write(", "), self.traverse, node.targets)
849
850 def visit_Assert(self, node):
851 self.fill("assert ")
852 self.traverse(node.test)
853 if node.msg:
854 self.write(", ")
855 self.traverse(node.msg)
856
857 def visit_Global(self, node):
858 self.fill("global ")
859 self.interleave(lambda: self.write(", "), self.write, node.names)
860
861 def visit_Nonlocal(self, node):
862 self.fill("nonlocal ")
863 self.interleave(lambda: self.write(", "), self.write, node.names)
864
865 def visit_Await(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300866 with self.require_parens(_Precedence.AWAIT, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300867 self.write("await")
868 if node.value:
869 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300870 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300871 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000872
873 def visit_Yield(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300874 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300875 self.write("yield")
876 if node.value:
877 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300878 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300879 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000880
881 def visit_YieldFrom(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300882 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300883 self.write("yield from ")
884 if not node.value:
885 raise ValueError("Node can't be used without a value attribute.")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300886 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300887 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000888
889 def visit_Raise(self, node):
890 self.fill("raise")
891 if not node.exc:
892 if node.cause:
893 raise ValueError(f"Node can't use cause without an exception.")
894 return
895 self.write(" ")
896 self.traverse(node.exc)
897 if node.cause:
898 self.write(" from ")
899 self.traverse(node.cause)
900
901 def visit_Try(self, node):
902 self.fill("try")
903 with self.block():
904 self.traverse(node.body)
905 for ex in node.handlers:
906 self.traverse(ex)
907 if node.orelse:
908 self.fill("else")
909 with self.block():
910 self.traverse(node.orelse)
911 if node.finalbody:
912 self.fill("finally")
913 with self.block():
914 self.traverse(node.finalbody)
915
916 def visit_ExceptHandler(self, node):
917 self.fill("except")
918 if node.type:
919 self.write(" ")
920 self.traverse(node.type)
921 if node.name:
922 self.write(" as ")
923 self.write(node.name)
924 with self.block():
925 self.traverse(node.body)
926
927 def visit_ClassDef(self, node):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300928 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000929 for deco in node.decorator_list:
930 self.fill("@")
931 self.traverse(deco)
932 self.fill("class " + node.name)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300933 with self.delimit("(", ")"):
934 comma = False
935 for e in node.bases:
936 if comma:
937 self.write(", ")
938 else:
939 comma = True
940 self.traverse(e)
941 for e in node.keywords:
942 if comma:
943 self.write(", ")
944 else:
945 comma = True
946 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000947
948 with self.block():
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300949 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000950
951 def visit_FunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300952 self._function_helper(node, "def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000953
954 def visit_AsyncFunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300955 self._function_helper(node, "async def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000956
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300957 def _function_helper(self, node, fill_suffix):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300958 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000959 for deco in node.decorator_list:
960 self.fill("@")
961 self.traverse(deco)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300962 def_str = fill_suffix + " " + node.name
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000963 self.fill(def_str)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300964 with self.delimit("(", ")"):
965 self.traverse(node.args)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000966 if node.returns:
967 self.write(" -> ")
968 self.traverse(node.returns)
969 with self.block():
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300970 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000971
972 def visit_For(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300973 self._for_helper("for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000974
975 def visit_AsyncFor(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300976 self._for_helper("async for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000977
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300978 def _for_helper(self, fill, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000979 self.fill(fill)
980 self.traverse(node.target)
981 self.write(" in ")
982 self.traverse(node.iter)
983 with self.block():
984 self.traverse(node.body)
985 if node.orelse:
986 self.fill("else")
987 with self.block():
988 self.traverse(node.orelse)
989
990 def visit_If(self, node):
991 self.fill("if ")
992 self.traverse(node.test)
993 with self.block():
994 self.traverse(node.body)
995 # collapse nested ifs into equivalent elifs.
996 while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], If):
997 node = node.orelse[0]
998 self.fill("elif ")
999 self.traverse(node.test)
1000 with self.block():
1001 self.traverse(node.body)
1002 # final else
1003 if node.orelse:
1004 self.fill("else")
1005 with self.block():
1006 self.traverse(node.orelse)
1007
1008 def visit_While(self, node):
1009 self.fill("while ")
1010 self.traverse(node.test)
1011 with self.block():
1012 self.traverse(node.body)
1013 if node.orelse:
1014 self.fill("else")
1015 with self.block():
1016 self.traverse(node.orelse)
1017
1018 def visit_With(self, node):
1019 self.fill("with ")
1020 self.interleave(lambda: self.write(", "), self.traverse, node.items)
1021 with self.block():
1022 self.traverse(node.body)
1023
1024 def visit_AsyncWith(self, node):
1025 self.fill("async with ")
1026 self.interleave(lambda: self.write(", "), self.traverse, node.items)
1027 with self.block():
1028 self.traverse(node.body)
1029
1030 def visit_JoinedStr(self, node):
1031 self.write("f")
1032 self._fstring_JoinedStr(node, self.buffer_writer)
1033 self.write(repr(self.buffer))
1034
1035 def visit_FormattedValue(self, node):
1036 self.write("f")
1037 self._fstring_FormattedValue(node, self.buffer_writer)
1038 self.write(repr(self.buffer))
1039
1040 def _fstring_JoinedStr(self, node, write):
1041 for value in node.values:
1042 meth = getattr(self, "_fstring_" + type(value).__name__)
1043 meth(value, write)
1044
1045 def _fstring_Constant(self, node, write):
1046 if not isinstance(node.value, str):
1047 raise ValueError("Constants inside JoinedStr should be a string.")
1048 value = node.value.replace("{", "{{").replace("}", "}}")
1049 write(value)
1050
1051 def _fstring_FormattedValue(self, node, write):
1052 write("{")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001053 unparser = type(self)()
1054 unparser.set_precedence(_Precedence.TEST.next(), node.value)
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +03001055 expr = unparser.visit(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001056 if expr.startswith("{"):
1057 write(" ") # Separate pair of opening brackets as "{ {"
1058 write(expr)
1059 if node.conversion != -1:
1060 conversion = chr(node.conversion)
1061 if conversion not in "sra":
1062 raise ValueError("Unknown f-string conversion.")
1063 write(f"!{conversion}")
1064 if node.format_spec:
1065 write(":")
1066 meth = getattr(self, "_fstring_" + type(node.format_spec).__name__)
1067 meth(node.format_spec, write)
1068 write("}")
1069
1070 def visit_Name(self, node):
1071 self.write(node.id)
1072
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001073 def _write_docstring(self, node):
1074 self.fill()
1075 if node.kind == "u":
1076 self.write("u")
1077
1078 # Preserve quotes in the docstring by escaping them
1079 value = node.value.replace("\\", "\\\\")
1080 value = value.replace('"""', '""\"')
1081 if value[-1] == '"':
1082 value = value.replace('"', '\\"', -1)
1083
1084 self.write(f'"""{value}"""')
1085
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001086 def _write_constant(self, value):
1087 if isinstance(value, (float, complex)):
1088 # Substitute overflowing decimal literal for AST infinities.
1089 self.write(repr(value).replace("inf", _INFSTR))
1090 else:
1091 self.write(repr(value))
1092
1093 def visit_Constant(self, node):
1094 value = node.value
1095 if isinstance(value, tuple):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001096 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001097 self.items_view(self._write_constant, value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001098 elif value is ...:
1099 self.write("...")
1100 else:
1101 if node.kind == "u":
1102 self.write("u")
1103 self._write_constant(node.value)
1104
1105 def visit_List(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001106 with self.delimit("[", "]"):
1107 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001108
1109 def visit_ListComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001110 with self.delimit("[", "]"):
1111 self.traverse(node.elt)
1112 for gen in node.generators:
1113 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001114
1115 def visit_GeneratorExp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001116 with self.delimit("(", ")"):
1117 self.traverse(node.elt)
1118 for gen in node.generators:
1119 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001120
1121 def visit_SetComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001122 with self.delimit("{", "}"):
1123 self.traverse(node.elt)
1124 for gen in node.generators:
1125 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001126
1127 def visit_DictComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001128 with self.delimit("{", "}"):
1129 self.traverse(node.key)
1130 self.write(": ")
1131 self.traverse(node.value)
1132 for gen in node.generators:
1133 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001134
1135 def visit_comprehension(self, node):
1136 if node.is_async:
1137 self.write(" async for ")
1138 else:
1139 self.write(" for ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001140 self.set_precedence(_Precedence.TUPLE, node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001141 self.traverse(node.target)
1142 self.write(" in ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001143 self.set_precedence(_Precedence.TEST.next(), node.iter, *node.ifs)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001144 self.traverse(node.iter)
1145 for if_clause in node.ifs:
1146 self.write(" if ")
1147 self.traverse(if_clause)
1148
1149 def visit_IfExp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001150 with self.require_parens(_Precedence.TEST, node):
1151 self.set_precedence(_Precedence.TEST.next(), node.body, node.test)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001152 self.traverse(node.body)
1153 self.write(" if ")
1154 self.traverse(node.test)
1155 self.write(" else ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001156 self.set_precedence(_Precedence.TEST, node.orelse)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001157 self.traverse(node.orelse)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001158
1159 def visit_Set(self, node):
1160 if not node.elts:
Shantanu01508dc2020-04-16 03:10:12 -07001161 raise ValueError("Set node should have at least one item")
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001162 with self.delimit("{", "}"):
1163 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001164
1165 def visit_Dict(self, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001166 def write_key_value_pair(k, v):
1167 self.traverse(k)
1168 self.write(": ")
1169 self.traverse(v)
1170
1171 def write_item(item):
1172 k, v = item
1173 if k is None:
1174 # for dictionary unpacking operator in dicts {**{'y': 2}}
1175 # see PEP 448 for details
1176 self.write("**")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001177 self.set_precedence(_Precedence.EXPR, v)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001178 self.traverse(v)
1179 else:
1180 write_key_value_pair(k, v)
1181
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001182 with self.delimit("{", "}"):
1183 self.interleave(
1184 lambda: self.write(", "), write_item, zip(node.keys, node.values)
1185 )
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001186
1187 def visit_Tuple(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001188 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001189 self.items_view(self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001190
1191 unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001192 unop_precedence = {
1193 "~": _Precedence.FACTOR,
1194 "not": _Precedence.NOT,
1195 "+": _Precedence.FACTOR,
1196 "-": _Precedence.FACTOR
1197 }
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001198
1199 def visit_UnaryOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001200 operator = self.unop[node.op.__class__.__name__]
1201 operator_precedence = self.unop_precedence[operator]
1202 with self.require_parens(operator_precedence, node):
1203 self.write(operator)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001204 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001205 self.set_precedence(operator_precedence, node.operand)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001206 self.traverse(node.operand)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001207
1208 binop = {
1209 "Add": "+",
1210 "Sub": "-",
1211 "Mult": "*",
1212 "MatMult": "@",
1213 "Div": "/",
1214 "Mod": "%",
1215 "LShift": "<<",
1216 "RShift": ">>",
1217 "BitOr": "|",
1218 "BitXor": "^",
1219 "BitAnd": "&",
1220 "FloorDiv": "//",
1221 "Pow": "**",
1222 }
1223
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001224 binop_precedence = {
1225 "+": _Precedence.ARITH,
1226 "-": _Precedence.ARITH,
1227 "*": _Precedence.TERM,
1228 "@": _Precedence.TERM,
1229 "/": _Precedence.TERM,
1230 "%": _Precedence.TERM,
1231 "<<": _Precedence.SHIFT,
1232 ">>": _Precedence.SHIFT,
1233 "|": _Precedence.BOR,
1234 "^": _Precedence.BXOR,
1235 "&": _Precedence.BAND,
1236 "//": _Precedence.TERM,
1237 "**": _Precedence.POWER,
1238 }
1239
1240 binop_rassoc = frozenset(("**",))
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001241 def visit_BinOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001242 operator = self.binop[node.op.__class__.__name__]
1243 operator_precedence = self.binop_precedence[operator]
1244 with self.require_parens(operator_precedence, node):
1245 if operator in self.binop_rassoc:
1246 left_precedence = operator_precedence.next()
1247 right_precedence = operator_precedence
1248 else:
1249 left_precedence = operator_precedence
1250 right_precedence = operator_precedence.next()
1251
1252 self.set_precedence(left_precedence, node.left)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001253 self.traverse(node.left)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001254 self.write(f" {operator} ")
1255 self.set_precedence(right_precedence, node.right)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001256 self.traverse(node.right)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001257
1258 cmpops = {
1259 "Eq": "==",
1260 "NotEq": "!=",
1261 "Lt": "<",
1262 "LtE": "<=",
1263 "Gt": ">",
1264 "GtE": ">=",
1265 "Is": "is",
1266 "IsNot": "is not",
1267 "In": "in",
1268 "NotIn": "not in",
1269 }
1270
1271 def visit_Compare(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001272 with self.require_parens(_Precedence.CMP, node):
1273 self.set_precedence(_Precedence.CMP.next(), node.left, *node.comparators)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001274 self.traverse(node.left)
1275 for o, e in zip(node.ops, node.comparators):
1276 self.write(" " + self.cmpops[o.__class__.__name__] + " ")
1277 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001278
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001279 boolops = {"And": "and", "Or": "or"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001280 boolop_precedence = {"and": _Precedence.AND, "or": _Precedence.OR}
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001281
1282 def visit_BoolOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001283 operator = self.boolops[node.op.__class__.__name__]
1284 operator_precedence = self.boolop_precedence[operator]
1285
1286 def increasing_level_traverse(node):
1287 nonlocal operator_precedence
1288 operator_precedence = operator_precedence.next()
1289 self.set_precedence(operator_precedence, node)
1290 self.traverse(node)
1291
1292 with self.require_parens(operator_precedence, node):
1293 s = f" {operator} "
1294 self.interleave(lambda: self.write(s), increasing_level_traverse, node.values)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001295
1296 def visit_Attribute(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001297 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001298 self.traverse(node.value)
1299 # Special case: 3.__abs__() is a syntax error, so if node.value
1300 # is an integer literal then we need to either parenthesize
1301 # it or add an extra space to get 3 .__abs__().
1302 if isinstance(node.value, Constant) and isinstance(node.value.value, int):
1303 self.write(" ")
1304 self.write(".")
1305 self.write(node.attr)
1306
1307 def visit_Call(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001308 self.set_precedence(_Precedence.ATOM, node.func)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001309 self.traverse(node.func)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001310 with self.delimit("(", ")"):
1311 comma = False
1312 for e in node.args:
1313 if comma:
1314 self.write(", ")
1315 else:
1316 comma = True
1317 self.traverse(e)
1318 for e in node.keywords:
1319 if comma:
1320 self.write(", ")
1321 else:
1322 comma = True
1323 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001324
1325 def visit_Subscript(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001326 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001327 self.traverse(node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001328 with self.delimit("[", "]"):
Serhiy Storchaka13d52c22020-03-10 18:52:34 +02001329 if isinstance(node.slice, Tuple) and node.slice.elts:
1330 self.items_view(self.traverse, node.slice.elts)
Serhiy Storchakac4928fc2020-03-07 17:25:32 +02001331 else:
1332 self.traverse(node.slice)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001333
1334 def visit_Starred(self, node):
1335 self.write("*")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001336 self.set_precedence(_Precedence.EXPR, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001337 self.traverse(node.value)
1338
1339 def visit_Ellipsis(self, node):
1340 self.write("...")
1341
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001342 def visit_Slice(self, node):
1343 if node.lower:
1344 self.traverse(node.lower)
1345 self.write(":")
1346 if node.upper:
1347 self.traverse(node.upper)
1348 if node.step:
1349 self.write(":")
1350 self.traverse(node.step)
1351
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001352 def visit_arg(self, node):
1353 self.write(node.arg)
1354 if node.annotation:
1355 self.write(": ")
1356 self.traverse(node.annotation)
1357
1358 def visit_arguments(self, node):
1359 first = True
1360 # normal arguments
1361 all_args = node.posonlyargs + node.args
1362 defaults = [None] * (len(all_args) - len(node.defaults)) + node.defaults
1363 for index, elements in enumerate(zip(all_args, defaults), 1):
1364 a, d = elements
1365 if first:
1366 first = False
1367 else:
1368 self.write(", ")
1369 self.traverse(a)
1370 if d:
1371 self.write("=")
1372 self.traverse(d)
1373 if index == len(node.posonlyargs):
1374 self.write(", /")
1375
1376 # varargs, or bare '*' if no varargs but keyword-only arguments present
1377 if node.vararg or node.kwonlyargs:
1378 if first:
1379 first = False
1380 else:
1381 self.write(", ")
1382 self.write("*")
1383 if node.vararg:
1384 self.write(node.vararg.arg)
1385 if node.vararg.annotation:
1386 self.write(": ")
1387 self.traverse(node.vararg.annotation)
1388
1389 # keyword-only arguments
1390 if node.kwonlyargs:
1391 for a, d in zip(node.kwonlyargs, node.kw_defaults):
Batuhan Taşkayaa322f502019-12-16 15:26:58 +03001392 self.write(", ")
1393 self.traverse(a)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001394 if d:
1395 self.write("=")
1396 self.traverse(d)
1397
1398 # kwargs
1399 if node.kwarg:
1400 if first:
1401 first = False
1402 else:
1403 self.write(", ")
1404 self.write("**" + node.kwarg.arg)
1405 if node.kwarg.annotation:
1406 self.write(": ")
1407 self.traverse(node.kwarg.annotation)
1408
1409 def visit_keyword(self, node):
1410 if node.arg is None:
1411 self.write("**")
1412 else:
1413 self.write(node.arg)
1414 self.write("=")
1415 self.traverse(node.value)
1416
1417 def visit_Lambda(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001418 with self.require_parens(_Precedence.TEST, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001419 self.write("lambda ")
1420 self.traverse(node.args)
1421 self.write(": ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001422 self.set_precedence(_Precedence.TEST, node.body)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001423 self.traverse(node.body)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001424
1425 def visit_alias(self, node):
1426 self.write(node.name)
1427 if node.asname:
1428 self.write(" as " + node.asname)
1429
1430 def visit_withitem(self, node):
1431 self.traverse(node.context_expr)
1432 if node.optional_vars:
1433 self.write(" as ")
1434 self.traverse(node.optional_vars)
1435
1436def unparse(ast_obj):
1437 unparser = _Unparser()
1438 return unparser.visit(ast_obj)
1439
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001440
1441def main():
1442 import argparse
1443
1444 parser = argparse.ArgumentParser(prog='python -m ast')
1445 parser.add_argument('infile', type=argparse.FileType(mode='rb'), nargs='?',
1446 default='-',
1447 help='the file to parse; defaults to stdin')
1448 parser.add_argument('-m', '--mode', default='exec',
1449 choices=('exec', 'single', 'eval', 'func_type'),
1450 help='specify what kind of code must be parsed')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001451 parser.add_argument('--no-type-comments', default=True, action='store_false',
1452 help="don't add information about type comments")
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001453 parser.add_argument('-a', '--include-attributes', action='store_true',
1454 help='include attributes such as line numbers and '
1455 'column offsets')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001456 parser.add_argument('-i', '--indent', type=int, default=3,
1457 help='indentation of nodes (number of spaces)')
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001458 args = parser.parse_args()
1459
1460 with args.infile as infile:
1461 source = infile.read()
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001462 tree = parse(source, args.infile.name, args.mode, type_comments=args.no_type_comments)
1463 print(dump(tree, include_attributes=args.include_attributes, indent=args.indent))
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001464
1465if __name__ == '__main__':
1466 main()