blob: 4839201e2e23463550bcc6fdb2d319c440ffe8cf [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
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020065 def _convert_num(node):
66 if isinstance(node, Constant):
Serhiy Storchaka3f228112018-09-27 17:42:37 +030067 if type(node.value) in (int, float, complex):
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020068 return node.value
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020069 raise ValueError('malformed node or string: ' + repr(node))
70 def _convert_signed_num(node):
71 if isinstance(node, UnaryOp) and isinstance(node.op, (UAdd, USub)):
72 operand = _convert_num(node.operand)
73 if isinstance(node.op, UAdd):
74 return + operand
75 else:
76 return - operand
77 return _convert_num(node)
Georg Brandl0c77a822008-06-10 16:37:50 +000078 def _convert(node):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010079 if isinstance(node, Constant):
80 return node.value
Georg Brandl0c77a822008-06-10 16:37:50 +000081 elif isinstance(node, Tuple):
82 return tuple(map(_convert, node.elts))
83 elif isinstance(node, List):
84 return list(map(_convert, node.elts))
Georg Brandl492f3fc2010-07-11 09:41:21 +000085 elif isinstance(node, Set):
86 return set(map(_convert, node.elts))
Raymond Hettinger4fcf5c12020-01-02 22:21:18 -070087 elif (isinstance(node, Call) and isinstance(node.func, Name) and
88 node.func.id == 'set' and node.args == node.keywords == []):
89 return set()
Georg Brandl0c77a822008-06-10 16:37:50 +000090 elif isinstance(node, Dict):
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020091 return dict(zip(map(_convert, node.keys),
92 map(_convert, node.values)))
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010093 elif isinstance(node, BinOp) and isinstance(node.op, (Add, Sub)):
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020094 left = _convert_signed_num(node.left)
95 right = _convert_num(node.right)
96 if isinstance(left, (int, float)) and isinstance(right, complex):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010097 if isinstance(node.op, Add):
98 return left + right
99 else:
100 return left - right
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +0200101 return _convert_signed_num(node)
Georg Brandl0c77a822008-06-10 16:37:50 +0000102 return _convert(node_or_string)
103
104
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300105def dump(node, annotate_fields=True, include_attributes=False, *, indent=None):
Georg Brandl0c77a822008-06-10 16:37:50 +0000106 """
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300107 Return a formatted dump of the tree in node. This is mainly useful for
108 debugging purposes. If annotate_fields is true (by default),
109 the returned string will show the names and the values for fields.
110 If annotate_fields is false, the result string will be more compact by
111 omitting unambiguous field names. Attributes such as line
Benjamin Petersondcf97b92008-07-02 17:30:14 +0000112 numbers and column offsets are not dumped by default. If this is wanted,
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300113 include_attributes can be set to true. If indent is a non-negative
114 integer or string, then the tree will be pretty-printed with that indent
115 level. None (the default) selects the single line representation.
Georg Brandl0c77a822008-06-10 16:37:50 +0000116 """
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300117 def _format(node, level=0):
118 if indent is not None:
119 level += 1
120 prefix = '\n' + indent * level
121 sep = ',\n' + indent * level
122 else:
123 prefix = ''
124 sep = ', '
Georg Brandl0c77a822008-06-10 16:37:50 +0000125 if isinstance(node, AST):
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300126 args = []
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300127 allsimple = True
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300128 keywords = annotate_fields
129 for field in node._fields:
130 try:
131 value = getattr(node, field)
132 except AttributeError:
133 keywords = True
134 else:
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300135 value, simple = _format(value, level)
136 allsimple = allsimple and simple
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300137 if keywords:
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300138 args.append('%s=%s' % (field, value))
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300139 else:
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300140 args.append(value)
Georg Brandl0c77a822008-06-10 16:37:50 +0000141 if include_attributes and node._attributes:
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300142 for attr in node._attributes:
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300143 try:
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300144 value = getattr(node, attr)
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300145 except AttributeError:
146 pass
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300147 else:
148 value, simple = _format(value, level)
149 allsimple = allsimple and simple
150 args.append('%s=%s' % (attr, value))
151 if allsimple and len(args) <= 3:
152 return '%s(%s)' % (node.__class__.__name__, ', '.join(args)), not args
153 return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args)), False
Georg Brandl0c77a822008-06-10 16:37:50 +0000154 elif isinstance(node, list):
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300155 if not node:
156 return '[]', True
157 return '[%s%s]' % (prefix, sep.join(_format(x, level)[0] for x in node)), False
158 return repr(node), True
159
Georg Brandl0c77a822008-06-10 16:37:50 +0000160 if not isinstance(node, AST):
161 raise TypeError('expected AST, got %r' % node.__class__.__name__)
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300162 if indent is not None and not isinstance(indent, str):
163 indent = ' ' * indent
164 return _format(node)[0]
Georg Brandl0c77a822008-06-10 16:37:50 +0000165
166
167def copy_location(new_node, old_node):
168 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000169 Copy source location (`lineno`, `col_offset`, `end_lineno`, and `end_col_offset`
170 attributes) from *old_node* to *new_node* if possible, and return *new_node*.
Georg Brandl0c77a822008-06-10 16:37:50 +0000171 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000172 for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
Georg Brandl0c77a822008-06-10 16:37:50 +0000173 if attr in old_node._attributes and attr in new_node._attributes \
174 and hasattr(old_node, attr):
175 setattr(new_node, attr, getattr(old_node, attr))
176 return new_node
177
178
179def fix_missing_locations(node):
180 """
181 When you compile a node tree with compile(), the compiler expects lineno and
182 col_offset attributes for every node that supports them. This is rather
183 tedious to fill in for generated nodes, so this helper adds these attributes
184 recursively where not already set, by setting them to the values of the
185 parent node. It works recursively starting at *node*.
186 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000187 def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
Georg Brandl0c77a822008-06-10 16:37:50 +0000188 if 'lineno' in node._attributes:
189 if not hasattr(node, 'lineno'):
190 node.lineno = lineno
191 else:
192 lineno = node.lineno
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000193 if 'end_lineno' in node._attributes:
194 if not hasattr(node, 'end_lineno'):
195 node.end_lineno = end_lineno
196 else:
197 end_lineno = node.end_lineno
Georg Brandl0c77a822008-06-10 16:37:50 +0000198 if 'col_offset' in node._attributes:
199 if not hasattr(node, 'col_offset'):
200 node.col_offset = col_offset
201 else:
202 col_offset = node.col_offset
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000203 if 'end_col_offset' in node._attributes:
204 if not hasattr(node, 'end_col_offset'):
205 node.end_col_offset = end_col_offset
206 else:
207 end_col_offset = node.end_col_offset
Georg Brandl0c77a822008-06-10 16:37:50 +0000208 for child in iter_child_nodes(node):
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000209 _fix(child, lineno, col_offset, end_lineno, end_col_offset)
210 _fix(node, 1, 0, 1, 0)
Georg Brandl0c77a822008-06-10 16:37:50 +0000211 return node
212
213
214def increment_lineno(node, n=1):
215 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000216 Increment the line number and end line number of each node in the tree
217 starting at *node* by *n*. This is useful to "move code" to a different
218 location in a file.
Georg Brandl0c77a822008-06-10 16:37:50 +0000219 """
Georg Brandl0c77a822008-06-10 16:37:50 +0000220 for child in walk(node):
221 if 'lineno' in child._attributes:
222 child.lineno = getattr(child, 'lineno', 0) + n
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000223 if 'end_lineno' in child._attributes:
224 child.end_lineno = getattr(child, 'end_lineno', 0) + n
Georg Brandl0c77a822008-06-10 16:37:50 +0000225 return node
226
227
228def iter_fields(node):
229 """
230 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
231 that is present on *node*.
232 """
233 for field in node._fields:
234 try:
235 yield field, getattr(node, field)
236 except AttributeError:
237 pass
238
239
240def iter_child_nodes(node):
241 """
242 Yield all direct child nodes of *node*, that is, all fields that are nodes
243 and all items of fields that are lists of nodes.
244 """
245 for name, field in iter_fields(node):
246 if isinstance(field, AST):
247 yield field
248 elif isinstance(field, list):
249 for item in field:
250 if isinstance(item, AST):
251 yield item
252
253
254def get_docstring(node, clean=True):
255 """
256 Return the docstring for the given node or None if no docstring can
257 be found. If the node provided does not have docstrings a TypeError
258 will be raised.
Matthias Bussonnier41cea702017-02-23 22:44:19 -0800259
260 If *clean* is `True`, all tabs are expanded to spaces and any whitespace
261 that can be uniformly removed from the second line onwards is removed.
Georg Brandl0c77a822008-06-10 16:37:50 +0000262 """
Yury Selivanov2f07a662015-07-23 08:54:35 +0300263 if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
Georg Brandl0c77a822008-06-10 16:37:50 +0000264 raise TypeError("%r can't have docstrings" % node.__class__.__name__)
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300265 if not(node.body and isinstance(node.body[0], Expr)):
Serhiy Storchaka73cbe7a2018-05-29 12:04:55 +0300266 return None
267 node = node.body[0].value
268 if isinstance(node, Str):
269 text = node.s
270 elif isinstance(node, Constant) and isinstance(node.value, str):
271 text = node.value
272 else:
273 return None
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300274 if clean:
Victor Stinnerf2c1aa12016-01-26 00:40:57 +0100275 import inspect
276 text = inspect.cleandoc(text)
277 return text
Georg Brandl0c77a822008-06-10 16:37:50 +0000278
279
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000280def _splitlines_no_ff(source):
281 """Split a string into lines ignoring form feed and other chars.
282
283 This mimics how the Python parser splits source code.
284 """
285 idx = 0
286 lines = []
287 next_line = ''
288 while idx < len(source):
289 c = source[idx]
290 next_line += c
291 idx += 1
292 # Keep \r\n together
293 if c == '\r' and idx < len(source) and source[idx] == '\n':
294 next_line += '\n'
295 idx += 1
296 if c in '\r\n':
297 lines.append(next_line)
298 next_line = ''
299
300 if next_line:
301 lines.append(next_line)
302 return lines
303
304
305def _pad_whitespace(source):
mpheathfbeba8f2020-02-14 04:32:09 +1000306 r"""Replace all chars except '\f\t' in a line with spaces."""
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000307 result = ''
308 for c in source:
309 if c in '\f\t':
310 result += c
311 else:
312 result += ' '
313 return result
314
315
316def get_source_segment(source, node, *, padded=False):
317 """Get source code segment of the *source* that generated *node*.
318
319 If some location information (`lineno`, `end_lineno`, `col_offset`,
320 or `end_col_offset`) is missing, return None.
321
322 If *padded* is `True`, the first line of a multi-line statement will
323 be padded with spaces to match its original position.
324 """
325 try:
326 lineno = node.lineno - 1
327 end_lineno = node.end_lineno - 1
328 col_offset = node.col_offset
329 end_col_offset = node.end_col_offset
330 except AttributeError:
331 return None
332
333 lines = _splitlines_no_ff(source)
334 if end_lineno == lineno:
335 return lines[lineno].encode()[col_offset:end_col_offset].decode()
336
337 if padded:
338 padding = _pad_whitespace(lines[lineno].encode()[:col_offset].decode())
339 else:
340 padding = ''
341
342 first = padding + lines[lineno].encode()[col_offset:].decode()
343 last = lines[end_lineno].encode()[:end_col_offset].decode()
344 lines = lines[lineno+1:end_lineno]
345
346 lines.insert(0, first)
347 lines.append(last)
348 return ''.join(lines)
349
350
Georg Brandl0c77a822008-06-10 16:37:50 +0000351def walk(node):
352 """
Georg Brandl619e7ba2011-01-09 07:38:51 +0000353 Recursively yield all descendant nodes in the tree starting at *node*
354 (including *node* itself), in no specified order. This is useful if you
355 only want to modify nodes in place and don't care about the context.
Georg Brandl0c77a822008-06-10 16:37:50 +0000356 """
357 from collections import deque
358 todo = deque([node])
359 while todo:
360 node = todo.popleft()
361 todo.extend(iter_child_nodes(node))
362 yield node
363
364
365class NodeVisitor(object):
366 """
367 A node visitor base class that walks the abstract syntax tree and calls a
368 visitor function for every node found. This function may return a value
369 which is forwarded by the `visit` method.
370
371 This class is meant to be subclassed, with the subclass adding visitor
372 methods.
373
374 Per default the visitor functions for the nodes are ``'visit_'`` +
375 class name of the node. So a `TryFinally` node visit function would
376 be `visit_TryFinally`. This behavior can be changed by overriding
377 the `visit` method. If no visitor function exists for a node
378 (return value `None`) the `generic_visit` visitor is used instead.
379
380 Don't use the `NodeVisitor` if you want to apply changes to nodes during
381 traversing. For this a special visitor exists (`NodeTransformer`) that
382 allows modifications.
383 """
384
385 def visit(self, node):
386 """Visit a node."""
387 method = 'visit_' + node.__class__.__name__
388 visitor = getattr(self, method, self.generic_visit)
389 return visitor(node)
390
391 def generic_visit(self, node):
392 """Called if no explicit visitor function exists for a node."""
393 for field, value in iter_fields(node):
394 if isinstance(value, list):
395 for item in value:
396 if isinstance(item, AST):
397 self.visit(item)
398 elif isinstance(value, AST):
399 self.visit(value)
400
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300401 def visit_Constant(self, node):
402 value = node.value
403 type_name = _const_node_type_names.get(type(value))
404 if type_name is None:
405 for cls, name in _const_node_type_names.items():
406 if isinstance(value, cls):
407 type_name = name
408 break
409 if type_name is not None:
410 method = 'visit_' + type_name
411 try:
412 visitor = getattr(self, method)
413 except AttributeError:
414 pass
415 else:
416 import warnings
417 warnings.warn(f"{method} is deprecated; add visit_Constant",
418 DeprecationWarning, 2)
419 return visitor(node)
420 return self.generic_visit(node)
421
Georg Brandl0c77a822008-06-10 16:37:50 +0000422
423class NodeTransformer(NodeVisitor):
424 """
425 A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
426 allows modification of nodes.
427
428 The `NodeTransformer` will walk the AST and use the return value of the
429 visitor methods to replace or remove the old node. If the return value of
430 the visitor method is ``None``, the node will be removed from its location,
431 otherwise it is replaced with the return value. The return value may be the
432 original node in which case no replacement takes place.
433
434 Here is an example transformer that rewrites all occurrences of name lookups
435 (``foo``) to ``data['foo']``::
436
437 class RewriteName(NodeTransformer):
438
439 def visit_Name(self, node):
440 return copy_location(Subscript(
441 value=Name(id='data', ctx=Load()),
442 slice=Index(value=Str(s=node.id)),
443 ctx=node.ctx
444 ), node)
445
446 Keep in mind that if the node you're operating on has child nodes you must
447 either transform the child nodes yourself or call the :meth:`generic_visit`
448 method for the node first.
449
450 For nodes that were part of a collection of statements (that applies to all
451 statement nodes), the visitor may also return a list of nodes rather than
452 just a single node.
453
454 Usually you use the transformer like this::
455
456 node = YourTransformer().visit(node)
457 """
458
459 def generic_visit(self, node):
460 for field, old_value in iter_fields(node):
Georg Brandl0c77a822008-06-10 16:37:50 +0000461 if isinstance(old_value, list):
462 new_values = []
463 for value in old_value:
464 if isinstance(value, AST):
465 value = self.visit(value)
466 if value is None:
467 continue
468 elif not isinstance(value, AST):
469 new_values.extend(value)
470 continue
471 new_values.append(value)
472 old_value[:] = new_values
473 elif isinstance(old_value, AST):
474 new_node = self.visit(old_value)
475 if new_node is None:
476 delattr(node, field)
477 else:
478 setattr(node, field, new_node)
479 return node
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300480
481
482# The following code is for backward compatibility.
483# It will be removed in future.
484
485def _getter(self):
486 return self.value
487
488def _setter(self, value):
489 self.value = value
490
491Constant.n = property(_getter, _setter)
492Constant.s = property(_getter, _setter)
493
494class _ABC(type):
495
496 def __instancecheck__(cls, inst):
497 if not isinstance(inst, Constant):
498 return False
499 if cls in _const_types:
500 try:
501 value = inst.value
502 except AttributeError:
503 return False
504 else:
Anthony Sottile74176222019-01-18 11:30:28 -0800505 return (
506 isinstance(value, _const_types[cls]) and
507 not isinstance(value, _const_types_not.get(cls, ()))
508 )
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300509 return type.__instancecheck__(cls, inst)
510
511def _new(cls, *args, **kwargs):
512 if cls in _const_types:
513 return Constant(*args, **kwargs)
514 return Constant.__new__(cls, *args, **kwargs)
515
516class Num(Constant, metaclass=_ABC):
517 _fields = ('n',)
518 __new__ = _new
519
520class Str(Constant, metaclass=_ABC):
521 _fields = ('s',)
522 __new__ = _new
523
524class Bytes(Constant, metaclass=_ABC):
525 _fields = ('s',)
526 __new__ = _new
527
528class NameConstant(Constant, metaclass=_ABC):
529 __new__ = _new
530
531class Ellipsis(Constant, metaclass=_ABC):
532 _fields = ()
533
534 def __new__(cls, *args, **kwargs):
535 if cls is Ellipsis:
536 return Constant(..., *args, **kwargs)
537 return Constant.__new__(cls, *args, **kwargs)
538
539_const_types = {
540 Num: (int, float, complex),
541 Str: (str,),
542 Bytes: (bytes,),
543 NameConstant: (type(None), bool),
544 Ellipsis: (type(...),),
545}
Anthony Sottile74176222019-01-18 11:30:28 -0800546_const_types_not = {
547 Num: (bool,),
548}
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300549_const_node_type_names = {
550 bool: 'NameConstant', # should be before int
551 type(None): 'NameConstant',
552 int: 'Num',
553 float: 'Num',
554 complex: 'Num',
555 str: 'Str',
556 bytes: 'Bytes',
557 type(...): 'Ellipsis',
558}
Serhiy Storchaka832e8642019-09-09 23:36:13 +0300559
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000560# Large float and imaginary literals get turned into infinities in the AST.
561# We unparse those infinities to INFSTR.
562_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
563
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300564class _Precedence(IntEnum):
565 """Precedence table that originated from python grammar."""
566
567 TUPLE = auto()
568 YIELD = auto() # 'yield', 'yield from'
569 TEST = auto() # 'if'-'else', 'lambda'
570 OR = auto() # 'or'
571 AND = auto() # 'and'
572 NOT = auto() # 'not'
573 CMP = auto() # '<', '>', '==', '>=', '<=', '!=',
574 # 'in', 'not in', 'is', 'is not'
575 EXPR = auto()
576 BOR = EXPR # '|'
577 BXOR = auto() # '^'
578 BAND = auto() # '&'
579 SHIFT = auto() # '<<', '>>'
580 ARITH = auto() # '+', '-'
581 TERM = auto() # '*', '@', '/', '%', '//'
582 FACTOR = auto() # unary '+', '-', '~'
583 POWER = auto() # '**'
584 AWAIT = auto() # 'await'
585 ATOM = auto()
586
587 def next(self):
588 try:
589 return self.__class__(self + 1)
590 except ValueError:
591 return self
592
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000593class _Unparser(NodeVisitor):
594 """Methods in this class recursively traverse an AST and
595 output source code for the abstract syntax; original formatting
596 is disregarded."""
597
598 def __init__(self):
599 self._source = []
600 self._buffer = []
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300601 self._precedences = {}
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000602 self._indent = 0
603
604 def interleave(self, inter, f, seq):
605 """Call f on each item in seq, calling inter() in between."""
606 seq = iter(seq)
607 try:
608 f(next(seq))
609 except StopIteration:
610 pass
611 else:
612 for x in seq:
613 inter()
614 f(x)
615
616 def fill(self, text=""):
617 """Indent a piece of text and append it, according to the current
618 indentation level"""
619 self.write("\n" + " " * self._indent + text)
620
621 def write(self, text):
622 """Append a piece of text"""
623 self._source.append(text)
624
625 def buffer_writer(self, text):
626 self._buffer.append(text)
627
628 @property
629 def buffer(self):
630 value = "".join(self._buffer)
631 self._buffer.clear()
632 return value
633
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000634 @contextmanager
635 def block(self):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000636 """A context manager for preparing the source for blocks. It adds
637 the character':', increases the indentation on enter and decreases
638 the indentation on exit."""
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000639 self.write(":")
640 self._indent += 1
641 yield
642 self._indent -= 1
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000643
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300644 @contextmanager
645 def delimit(self, start, end):
646 """A context manager for preparing the source for expressions. It adds
647 *start* to the buffer and enters, after exit it adds *end*."""
648
649 self.write(start)
650 yield
651 self.write(end)
652
653 def delimit_if(self, start, end, condition):
654 if condition:
655 return self.delimit(start, end)
656 else:
657 return nullcontext()
658
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300659 def require_parens(self, precedence, node):
660 """Shortcut to adding precedence related parens"""
661 return self.delimit_if("(", ")", self.get_precedence(node) > precedence)
662
663 def get_precedence(self, node):
664 return self._precedences.get(node, _Precedence.TEST)
665
666 def set_precedence(self, precedence, *nodes):
667 for node in nodes:
668 self._precedences[node] = precedence
669
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000670 def traverse(self, node):
671 if isinstance(node, list):
672 for item in node:
673 self.traverse(item)
674 else:
675 super().visit(node)
676
677 def visit(self, node):
678 """Outputs a source code string that, if converted back to an ast
679 (using ast.parse) will generate an AST equivalent to *node*"""
680 self._source = []
681 self.traverse(node)
682 return "".join(self._source)
683
684 def visit_Module(self, node):
685 for subnode in node.body:
686 self.traverse(subnode)
687
688 def visit_Expr(self, node):
689 self.fill()
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300690 self.set_precedence(_Precedence.YIELD, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000691 self.traverse(node.value)
692
693 def visit_NamedExpr(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300694 with self.require_parens(_Precedence.TUPLE, node):
695 self.set_precedence(_Precedence.ATOM, node.target, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300696 self.traverse(node.target)
697 self.write(" := ")
698 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000699
700 def visit_Import(self, node):
701 self.fill("import ")
702 self.interleave(lambda: self.write(", "), self.traverse, node.names)
703
704 def visit_ImportFrom(self, node):
705 self.fill("from ")
706 self.write("." * node.level)
707 if node.module:
708 self.write(node.module)
709 self.write(" import ")
710 self.interleave(lambda: self.write(", "), self.traverse, node.names)
711
712 def visit_Assign(self, node):
713 self.fill()
714 for target in node.targets:
715 self.traverse(target)
716 self.write(" = ")
717 self.traverse(node.value)
718
719 def visit_AugAssign(self, node):
720 self.fill()
721 self.traverse(node.target)
722 self.write(" " + self.binop[node.op.__class__.__name__] + "= ")
723 self.traverse(node.value)
724
725 def visit_AnnAssign(self, node):
726 self.fill()
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300727 with self.delimit_if("(", ")", not node.simple and isinstance(node.target, Name)):
728 self.traverse(node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000729 self.write(": ")
730 self.traverse(node.annotation)
731 if node.value:
732 self.write(" = ")
733 self.traverse(node.value)
734
735 def visit_Return(self, node):
736 self.fill("return")
737 if node.value:
738 self.write(" ")
739 self.traverse(node.value)
740
741 def visit_Pass(self, node):
742 self.fill("pass")
743
744 def visit_Break(self, node):
745 self.fill("break")
746
747 def visit_Continue(self, node):
748 self.fill("continue")
749
750 def visit_Delete(self, node):
751 self.fill("del ")
752 self.interleave(lambda: self.write(", "), self.traverse, node.targets)
753
754 def visit_Assert(self, node):
755 self.fill("assert ")
756 self.traverse(node.test)
757 if node.msg:
758 self.write(", ")
759 self.traverse(node.msg)
760
761 def visit_Global(self, node):
762 self.fill("global ")
763 self.interleave(lambda: self.write(", "), self.write, node.names)
764
765 def visit_Nonlocal(self, node):
766 self.fill("nonlocal ")
767 self.interleave(lambda: self.write(", "), self.write, node.names)
768
769 def visit_Await(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300770 with self.require_parens(_Precedence.AWAIT, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300771 self.write("await")
772 if node.value:
773 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300774 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300775 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000776
777 def visit_Yield(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300778 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300779 self.write("yield")
780 if node.value:
781 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300782 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300783 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000784
785 def visit_YieldFrom(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300786 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300787 self.write("yield from ")
788 if not node.value:
789 raise ValueError("Node can't be used without a value attribute.")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300790 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300791 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000792
793 def visit_Raise(self, node):
794 self.fill("raise")
795 if not node.exc:
796 if node.cause:
797 raise ValueError(f"Node can't use cause without an exception.")
798 return
799 self.write(" ")
800 self.traverse(node.exc)
801 if node.cause:
802 self.write(" from ")
803 self.traverse(node.cause)
804
805 def visit_Try(self, node):
806 self.fill("try")
807 with self.block():
808 self.traverse(node.body)
809 for ex in node.handlers:
810 self.traverse(ex)
811 if node.orelse:
812 self.fill("else")
813 with self.block():
814 self.traverse(node.orelse)
815 if node.finalbody:
816 self.fill("finally")
817 with self.block():
818 self.traverse(node.finalbody)
819
820 def visit_ExceptHandler(self, node):
821 self.fill("except")
822 if node.type:
823 self.write(" ")
824 self.traverse(node.type)
825 if node.name:
826 self.write(" as ")
827 self.write(node.name)
828 with self.block():
829 self.traverse(node.body)
830
831 def visit_ClassDef(self, node):
832 self.write("\n")
833 for deco in node.decorator_list:
834 self.fill("@")
835 self.traverse(deco)
836 self.fill("class " + node.name)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300837 with self.delimit("(", ")"):
838 comma = False
839 for e in node.bases:
840 if comma:
841 self.write(", ")
842 else:
843 comma = True
844 self.traverse(e)
845 for e in node.keywords:
846 if comma:
847 self.write(", ")
848 else:
849 comma = True
850 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000851
852 with self.block():
853 self.traverse(node.body)
854
855 def visit_FunctionDef(self, node):
856 self.__FunctionDef_helper(node, "def")
857
858 def visit_AsyncFunctionDef(self, node):
859 self.__FunctionDef_helper(node, "async def")
860
861 def __FunctionDef_helper(self, node, fill_suffix):
862 self.write("\n")
863 for deco in node.decorator_list:
864 self.fill("@")
865 self.traverse(deco)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300866 def_str = fill_suffix + " " + node.name
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000867 self.fill(def_str)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300868 with self.delimit("(", ")"):
869 self.traverse(node.args)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000870 if node.returns:
871 self.write(" -> ")
872 self.traverse(node.returns)
873 with self.block():
874 self.traverse(node.body)
875
876 def visit_For(self, node):
877 self.__For_helper("for ", node)
878
879 def visit_AsyncFor(self, node):
880 self.__For_helper("async for ", node)
881
882 def __For_helper(self, fill, node):
883 self.fill(fill)
884 self.traverse(node.target)
885 self.write(" in ")
886 self.traverse(node.iter)
887 with self.block():
888 self.traverse(node.body)
889 if node.orelse:
890 self.fill("else")
891 with self.block():
892 self.traverse(node.orelse)
893
894 def visit_If(self, node):
895 self.fill("if ")
896 self.traverse(node.test)
897 with self.block():
898 self.traverse(node.body)
899 # collapse nested ifs into equivalent elifs.
900 while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], If):
901 node = node.orelse[0]
902 self.fill("elif ")
903 self.traverse(node.test)
904 with self.block():
905 self.traverse(node.body)
906 # final else
907 if node.orelse:
908 self.fill("else")
909 with self.block():
910 self.traverse(node.orelse)
911
912 def visit_While(self, node):
913 self.fill("while ")
914 self.traverse(node.test)
915 with self.block():
916 self.traverse(node.body)
917 if node.orelse:
918 self.fill("else")
919 with self.block():
920 self.traverse(node.orelse)
921
922 def visit_With(self, node):
923 self.fill("with ")
924 self.interleave(lambda: self.write(", "), self.traverse, node.items)
925 with self.block():
926 self.traverse(node.body)
927
928 def visit_AsyncWith(self, node):
929 self.fill("async with ")
930 self.interleave(lambda: self.write(", "), self.traverse, node.items)
931 with self.block():
932 self.traverse(node.body)
933
934 def visit_JoinedStr(self, node):
935 self.write("f")
936 self._fstring_JoinedStr(node, self.buffer_writer)
937 self.write(repr(self.buffer))
938
939 def visit_FormattedValue(self, node):
940 self.write("f")
941 self._fstring_FormattedValue(node, self.buffer_writer)
942 self.write(repr(self.buffer))
943
944 def _fstring_JoinedStr(self, node, write):
945 for value in node.values:
946 meth = getattr(self, "_fstring_" + type(value).__name__)
947 meth(value, write)
948
949 def _fstring_Constant(self, node, write):
950 if not isinstance(node.value, str):
951 raise ValueError("Constants inside JoinedStr should be a string.")
952 value = node.value.replace("{", "{{").replace("}", "}}")
953 write(value)
954
955 def _fstring_FormattedValue(self, node, write):
956 write("{")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300957 unparser = type(self)()
958 unparser.set_precedence(_Precedence.TEST.next(), node.value)
959 expr = unparser.visit(node.value).rstrip("\n")
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000960 if expr.startswith("{"):
961 write(" ") # Separate pair of opening brackets as "{ {"
962 write(expr)
963 if node.conversion != -1:
964 conversion = chr(node.conversion)
965 if conversion not in "sra":
966 raise ValueError("Unknown f-string conversion.")
967 write(f"!{conversion}")
968 if node.format_spec:
969 write(":")
970 meth = getattr(self, "_fstring_" + type(node.format_spec).__name__)
971 meth(node.format_spec, write)
972 write("}")
973
974 def visit_Name(self, node):
975 self.write(node.id)
976
977 def _write_constant(self, value):
978 if isinstance(value, (float, complex)):
979 # Substitute overflowing decimal literal for AST infinities.
980 self.write(repr(value).replace("inf", _INFSTR))
981 else:
982 self.write(repr(value))
983
984 def visit_Constant(self, node):
985 value = node.value
986 if isinstance(value, tuple):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300987 with self.delimit("(", ")"):
988 if len(value) == 1:
989 self._write_constant(value[0])
990 self.write(",")
991 else:
992 self.interleave(lambda: self.write(", "), self._write_constant, value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000993 elif value is ...:
994 self.write("...")
995 else:
996 if node.kind == "u":
997 self.write("u")
998 self._write_constant(node.value)
999
1000 def visit_List(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001001 with self.delimit("[", "]"):
1002 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001003
1004 def visit_ListComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001005 with self.delimit("[", "]"):
1006 self.traverse(node.elt)
1007 for gen in node.generators:
1008 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001009
1010 def visit_GeneratorExp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001011 with self.delimit("(", ")"):
1012 self.traverse(node.elt)
1013 for gen in node.generators:
1014 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001015
1016 def visit_SetComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001017 with self.delimit("{", "}"):
1018 self.traverse(node.elt)
1019 for gen in node.generators:
1020 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001021
1022 def visit_DictComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001023 with self.delimit("{", "}"):
1024 self.traverse(node.key)
1025 self.write(": ")
1026 self.traverse(node.value)
1027 for gen in node.generators:
1028 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001029
1030 def visit_comprehension(self, node):
1031 if node.is_async:
1032 self.write(" async for ")
1033 else:
1034 self.write(" for ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001035 self.set_precedence(_Precedence.TUPLE, node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001036 self.traverse(node.target)
1037 self.write(" in ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001038 self.set_precedence(_Precedence.TEST.next(), node.iter, *node.ifs)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001039 self.traverse(node.iter)
1040 for if_clause in node.ifs:
1041 self.write(" if ")
1042 self.traverse(if_clause)
1043
1044 def visit_IfExp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001045 with self.require_parens(_Precedence.TEST, node):
1046 self.set_precedence(_Precedence.TEST.next(), node.body, node.test)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001047 self.traverse(node.body)
1048 self.write(" if ")
1049 self.traverse(node.test)
1050 self.write(" else ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001051 self.set_precedence(_Precedence.TEST, node.orelse)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001052 self.traverse(node.orelse)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001053
1054 def visit_Set(self, node):
1055 if not node.elts:
1056 raise ValueError("Set node should has at least one item")
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001057 with self.delimit("{", "}"):
1058 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001059
1060 def visit_Dict(self, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001061 def write_key_value_pair(k, v):
1062 self.traverse(k)
1063 self.write(": ")
1064 self.traverse(v)
1065
1066 def write_item(item):
1067 k, v = item
1068 if k is None:
1069 # for dictionary unpacking operator in dicts {**{'y': 2}}
1070 # see PEP 448 for details
1071 self.write("**")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001072 self.set_precedence(_Precedence.EXPR, v)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001073 self.traverse(v)
1074 else:
1075 write_key_value_pair(k, v)
1076
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001077 with self.delimit("{", "}"):
1078 self.interleave(
1079 lambda: self.write(", "), write_item, zip(node.keys, node.values)
1080 )
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001081
1082 def visit_Tuple(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001083 with self.delimit("(", ")"):
1084 if len(node.elts) == 1:
1085 elt = node.elts[0]
1086 self.traverse(elt)
1087 self.write(",")
1088 else:
1089 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001090
1091 unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001092 unop_precedence = {
1093 "~": _Precedence.FACTOR,
1094 "not": _Precedence.NOT,
1095 "+": _Precedence.FACTOR,
1096 "-": _Precedence.FACTOR
1097 }
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001098
1099 def visit_UnaryOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001100 operator = self.unop[node.op.__class__.__name__]
1101 operator_precedence = self.unop_precedence[operator]
1102 with self.require_parens(operator_precedence, node):
1103 self.write(operator)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001104 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001105 self.set_precedence(operator_precedence, node.operand)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001106 self.traverse(node.operand)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001107
1108 binop = {
1109 "Add": "+",
1110 "Sub": "-",
1111 "Mult": "*",
1112 "MatMult": "@",
1113 "Div": "/",
1114 "Mod": "%",
1115 "LShift": "<<",
1116 "RShift": ">>",
1117 "BitOr": "|",
1118 "BitXor": "^",
1119 "BitAnd": "&",
1120 "FloorDiv": "//",
1121 "Pow": "**",
1122 }
1123
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001124 binop_precedence = {
1125 "+": _Precedence.ARITH,
1126 "-": _Precedence.ARITH,
1127 "*": _Precedence.TERM,
1128 "@": _Precedence.TERM,
1129 "/": _Precedence.TERM,
1130 "%": _Precedence.TERM,
1131 "<<": _Precedence.SHIFT,
1132 ">>": _Precedence.SHIFT,
1133 "|": _Precedence.BOR,
1134 "^": _Precedence.BXOR,
1135 "&": _Precedence.BAND,
1136 "//": _Precedence.TERM,
1137 "**": _Precedence.POWER,
1138 }
1139
1140 binop_rassoc = frozenset(("**",))
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001141 def visit_BinOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001142 operator = self.binop[node.op.__class__.__name__]
1143 operator_precedence = self.binop_precedence[operator]
1144 with self.require_parens(operator_precedence, node):
1145 if operator in self.binop_rassoc:
1146 left_precedence = operator_precedence.next()
1147 right_precedence = operator_precedence
1148 else:
1149 left_precedence = operator_precedence
1150 right_precedence = operator_precedence.next()
1151
1152 self.set_precedence(left_precedence, node.left)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001153 self.traverse(node.left)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001154 self.write(f" {operator} ")
1155 self.set_precedence(right_precedence, node.right)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001156 self.traverse(node.right)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001157
1158 cmpops = {
1159 "Eq": "==",
1160 "NotEq": "!=",
1161 "Lt": "<",
1162 "LtE": "<=",
1163 "Gt": ">",
1164 "GtE": ">=",
1165 "Is": "is",
1166 "IsNot": "is not",
1167 "In": "in",
1168 "NotIn": "not in",
1169 }
1170
1171 def visit_Compare(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001172 with self.require_parens(_Precedence.CMP, node):
1173 self.set_precedence(_Precedence.CMP.next(), node.left, *node.comparators)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001174 self.traverse(node.left)
1175 for o, e in zip(node.ops, node.comparators):
1176 self.write(" " + self.cmpops[o.__class__.__name__] + " ")
1177 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001178
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001179 boolops = {"And": "and", "Or": "or"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001180 boolop_precedence = {"and": _Precedence.AND, "or": _Precedence.OR}
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001181
1182 def visit_BoolOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001183 operator = self.boolops[node.op.__class__.__name__]
1184 operator_precedence = self.boolop_precedence[operator]
1185
1186 def increasing_level_traverse(node):
1187 nonlocal operator_precedence
1188 operator_precedence = operator_precedence.next()
1189 self.set_precedence(operator_precedence, node)
1190 self.traverse(node)
1191
1192 with self.require_parens(operator_precedence, node):
1193 s = f" {operator} "
1194 self.interleave(lambda: self.write(s), increasing_level_traverse, node.values)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001195
1196 def visit_Attribute(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001197 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001198 self.traverse(node.value)
1199 # Special case: 3.__abs__() is a syntax error, so if node.value
1200 # is an integer literal then we need to either parenthesize
1201 # it or add an extra space to get 3 .__abs__().
1202 if isinstance(node.value, Constant) and isinstance(node.value.value, int):
1203 self.write(" ")
1204 self.write(".")
1205 self.write(node.attr)
1206
1207 def visit_Call(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001208 self.set_precedence(_Precedence.ATOM, node.func)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001209 self.traverse(node.func)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001210 with self.delimit("(", ")"):
1211 comma = False
1212 for e in node.args:
1213 if comma:
1214 self.write(", ")
1215 else:
1216 comma = True
1217 self.traverse(e)
1218 for e in node.keywords:
1219 if comma:
1220 self.write(", ")
1221 else:
1222 comma = True
1223 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001224
1225 def visit_Subscript(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001226 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001227 self.traverse(node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001228 with self.delimit("[", "]"):
1229 self.traverse(node.slice)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001230
1231 def visit_Starred(self, node):
1232 self.write("*")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001233 self.set_precedence(_Precedence.EXPR, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001234 self.traverse(node.value)
1235
1236 def visit_Ellipsis(self, node):
1237 self.write("...")
1238
1239 def visit_Index(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001240 self.set_precedence(_Precedence.TUPLE, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001241 self.traverse(node.value)
1242
1243 def visit_Slice(self, node):
1244 if node.lower:
1245 self.traverse(node.lower)
1246 self.write(":")
1247 if node.upper:
1248 self.traverse(node.upper)
1249 if node.step:
1250 self.write(":")
1251 self.traverse(node.step)
1252
1253 def visit_ExtSlice(self, node):
1254 self.interleave(lambda: self.write(", "), self.traverse, node.dims)
1255
1256 def visit_arg(self, node):
1257 self.write(node.arg)
1258 if node.annotation:
1259 self.write(": ")
1260 self.traverse(node.annotation)
1261
1262 def visit_arguments(self, node):
1263 first = True
1264 # normal arguments
1265 all_args = node.posonlyargs + node.args
1266 defaults = [None] * (len(all_args) - len(node.defaults)) + node.defaults
1267 for index, elements in enumerate(zip(all_args, defaults), 1):
1268 a, d = elements
1269 if first:
1270 first = False
1271 else:
1272 self.write(", ")
1273 self.traverse(a)
1274 if d:
1275 self.write("=")
1276 self.traverse(d)
1277 if index == len(node.posonlyargs):
1278 self.write(", /")
1279
1280 # varargs, or bare '*' if no varargs but keyword-only arguments present
1281 if node.vararg or node.kwonlyargs:
1282 if first:
1283 first = False
1284 else:
1285 self.write(", ")
1286 self.write("*")
1287 if node.vararg:
1288 self.write(node.vararg.arg)
1289 if node.vararg.annotation:
1290 self.write(": ")
1291 self.traverse(node.vararg.annotation)
1292
1293 # keyword-only arguments
1294 if node.kwonlyargs:
1295 for a, d in zip(node.kwonlyargs, node.kw_defaults):
Batuhan Taşkayaa322f502019-12-16 15:26:58 +03001296 self.write(", ")
1297 self.traverse(a)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001298 if d:
1299 self.write("=")
1300 self.traverse(d)
1301
1302 # kwargs
1303 if node.kwarg:
1304 if first:
1305 first = False
1306 else:
1307 self.write(", ")
1308 self.write("**" + node.kwarg.arg)
1309 if node.kwarg.annotation:
1310 self.write(": ")
1311 self.traverse(node.kwarg.annotation)
1312
1313 def visit_keyword(self, node):
1314 if node.arg is None:
1315 self.write("**")
1316 else:
1317 self.write(node.arg)
1318 self.write("=")
1319 self.traverse(node.value)
1320
1321 def visit_Lambda(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001322 with self.require_parens(_Precedence.TEST, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001323 self.write("lambda ")
1324 self.traverse(node.args)
1325 self.write(": ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001326 self.set_precedence(_Precedence.TEST, node.body)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001327 self.traverse(node.body)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001328
1329 def visit_alias(self, node):
1330 self.write(node.name)
1331 if node.asname:
1332 self.write(" as " + node.asname)
1333
1334 def visit_withitem(self, node):
1335 self.traverse(node.context_expr)
1336 if node.optional_vars:
1337 self.write(" as ")
1338 self.traverse(node.optional_vars)
1339
1340def unparse(ast_obj):
1341 unparser = _Unparser()
1342 return unparser.visit(ast_obj)
1343
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001344
1345def main():
1346 import argparse
1347
1348 parser = argparse.ArgumentParser(prog='python -m ast')
1349 parser.add_argument('infile', type=argparse.FileType(mode='rb'), nargs='?',
1350 default='-',
1351 help='the file to parse; defaults to stdin')
1352 parser.add_argument('-m', '--mode', default='exec',
1353 choices=('exec', 'single', 'eval', 'func_type'),
1354 help='specify what kind of code must be parsed')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001355 parser.add_argument('--no-type-comments', default=True, action='store_false',
1356 help="don't add information about type comments")
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001357 parser.add_argument('-a', '--include-attributes', action='store_true',
1358 help='include attributes such as line numbers and '
1359 'column offsets')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001360 parser.add_argument('-i', '--indent', type=int, default=3,
1361 help='indentation of nodes (number of spaces)')
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001362 args = parser.parse_args()
1363
1364 with args.infile as infile:
1365 source = infile.read()
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001366 tree = parse(source, args.infile.name, args.mode, type_comments=args.no_type_comments)
1367 print(dump(tree, include_attributes=args.include_attributes, indent=args.indent))
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001368
1369if __name__ == '__main__':
1370 main()