blob: f4d2f6e42c022be5fad3167206c26758f99ee844 [file] [log] [blame]
Haibo Huangd8830302020-03-03 10:09:46 -08001"""
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"""
Haibo Huang5eba2b42021-01-22 11:22:02 -080027import sys
Haibo Huangd8830302020-03-03 10:09:46 -080028from _ast import *
Haibo Huang5eba2b42021-01-22 11:22:02 -080029from contextlib import contextmanager, nullcontext
30from enum import IntEnum, auto
Haibo Huangd8830302020-03-03 10:09:46 -080031
32
33def parse(source, filename='<unknown>', mode='exec', *,
34 type_comments=False, feature_version=None):
35 """
36 Parse the source into an AST node.
37 Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
38 Pass type_comments=True to get back type comments where the syntax allows.
39 """
40 flags = PyCF_ONLY_AST
41 if type_comments:
42 flags |= PyCF_TYPE_COMMENTS
43 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.
50 return compile(source, filename, mode, flags,
51 _feature_version=feature_version)
52
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
58 Python literal structures: strings, bytes, numbers, tuples, lists, dicts,
59 sets, booleans, and None.
60 """
61 if isinstance(node_or_string, str):
Yi Kong71199322022-08-30 15:53:45 +080062 node_or_string = parse(node_or_string.lstrip(" \t"), mode='eval')
Haibo Huangd8830302020-03-03 10:09:46 -080063 if isinstance(node_or_string, Expression):
64 node_or_string = node_or_string.body
Haibo Huangf5f93a72020-10-19 15:43:42 -070065 def _raise_malformed_node(node):
Yi Kong71199322022-08-30 15:53:45 +080066 msg = "malformed node or string"
67 if lno := getattr(node, 'lineno', None):
68 msg += f' on line {lno}'
69 raise ValueError(msg + f': {node!r}')
Haibo Huangd8830302020-03-03 10:09:46 -080070 def _convert_num(node):
Haibo Huangf5f93a72020-10-19 15:43:42 -070071 if not isinstance(node, Constant) or type(node.value) not in (int, float, complex):
72 _raise_malformed_node(node)
73 return node.value
Haibo Huangd8830302020-03-03 10:09:46 -080074 def _convert_signed_num(node):
75 if isinstance(node, UnaryOp) and isinstance(node.op, (UAdd, USub)):
76 operand = _convert_num(node.operand)
77 if isinstance(node.op, UAdd):
78 return + operand
79 else:
80 return - operand
81 return _convert_num(node)
82 def _convert(node):
83 if isinstance(node, Constant):
84 return node.value
85 elif isinstance(node, Tuple):
86 return tuple(map(_convert, node.elts))
87 elif isinstance(node, List):
88 return list(map(_convert, node.elts))
89 elif isinstance(node, Set):
90 return set(map(_convert, node.elts))
Haibo Huang5eba2b42021-01-22 11:22:02 -080091 elif (isinstance(node, Call) and isinstance(node.func, Name) and
92 node.func.id == 'set' and node.args == node.keywords == []):
93 return set()
Haibo Huangd8830302020-03-03 10:09:46 -080094 elif isinstance(node, Dict):
Haibo Huangf5f93a72020-10-19 15:43:42 -070095 if len(node.keys) != len(node.values):
96 _raise_malformed_node(node)
Haibo Huangd8830302020-03-03 10:09:46 -080097 return dict(zip(map(_convert, node.keys),
98 map(_convert, node.values)))
99 elif isinstance(node, BinOp) and isinstance(node.op, (Add, Sub)):
100 left = _convert_signed_num(node.left)
101 right = _convert_num(node.right)
102 if isinstance(left, (int, float)) and isinstance(right, complex):
103 if isinstance(node.op, Add):
104 return left + right
105 else:
106 return left - right
107 return _convert_signed_num(node)
108 return _convert(node_or_string)
109
110
Haibo Huang5eba2b42021-01-22 11:22:02 -0800111def dump(node, annotate_fields=True, include_attributes=False, *, indent=None):
Haibo Huangd8830302020-03-03 10:09:46 -0800112 """
113 Return a formatted dump of the tree in node. This is mainly useful for
114 debugging purposes. If annotate_fields is true (by default),
115 the returned string will show the names and the values for fields.
116 If annotate_fields is false, the result string will be more compact by
117 omitting unambiguous field names. Attributes such as line
118 numbers and column offsets are not dumped by default. If this is wanted,
Haibo Huang5eba2b42021-01-22 11:22:02 -0800119 include_attributes can be set to true. If indent is a non-negative
120 integer or string, then the tree will be pretty-printed with that indent
121 level. None (the default) selects the single line representation.
Haibo Huangd8830302020-03-03 10:09:46 -0800122 """
Haibo Huang5eba2b42021-01-22 11:22:02 -0800123 def _format(node, level=0):
124 if indent is not None:
125 level += 1
126 prefix = '\n' + indent * level
127 sep = ',\n' + indent * level
128 else:
129 prefix = ''
130 sep = ', '
Haibo Huangd8830302020-03-03 10:09:46 -0800131 if isinstance(node, AST):
Haibo Huang5eba2b42021-01-22 11:22:02 -0800132 cls = type(node)
Haibo Huangd8830302020-03-03 10:09:46 -0800133 args = []
Haibo Huang5eba2b42021-01-22 11:22:02 -0800134 allsimple = True
Haibo Huangd8830302020-03-03 10:09:46 -0800135 keywords = annotate_fields
Haibo Huang5eba2b42021-01-22 11:22:02 -0800136 for name in node._fields:
Haibo Huangd8830302020-03-03 10:09:46 -0800137 try:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800138 value = getattr(node, name)
Haibo Huangd8830302020-03-03 10:09:46 -0800139 except AttributeError:
140 keywords = True
Haibo Huang5eba2b42021-01-22 11:22:02 -0800141 continue
142 if value is None and getattr(cls, name, ...) is None:
143 keywords = True
144 continue
145 value, simple = _format(value, level)
146 allsimple = allsimple and simple
147 if keywords:
148 args.append('%s=%s' % (name, value))
Haibo Huangd8830302020-03-03 10:09:46 -0800149 else:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800150 args.append(value)
Haibo Huangd8830302020-03-03 10:09:46 -0800151 if include_attributes and node._attributes:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800152 for name in node._attributes:
Haibo Huangd8830302020-03-03 10:09:46 -0800153 try:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800154 value = getattr(node, name)
Haibo Huangd8830302020-03-03 10:09:46 -0800155 except AttributeError:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800156 continue
157 if value is None and getattr(cls, name, ...) is None:
158 continue
159 value, simple = _format(value, level)
160 allsimple = allsimple and simple
161 args.append('%s=%s' % (name, value))
162 if allsimple and len(args) <= 3:
163 return '%s(%s)' % (node.__class__.__name__, ', '.join(args)), not args
164 return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args)), False
Haibo Huangd8830302020-03-03 10:09:46 -0800165 elif isinstance(node, list):
Haibo Huang5eba2b42021-01-22 11:22:02 -0800166 if not node:
167 return '[]', True
168 return '[%s%s]' % (prefix, sep.join(_format(x, level)[0] for x in node)), False
169 return repr(node), True
170
Haibo Huangd8830302020-03-03 10:09:46 -0800171 if not isinstance(node, AST):
172 raise TypeError('expected AST, got %r' % node.__class__.__name__)
Haibo Huang5eba2b42021-01-22 11:22:02 -0800173 if indent is not None and not isinstance(indent, str):
174 indent = ' ' * indent
175 return _format(node)[0]
Haibo Huangd8830302020-03-03 10:09:46 -0800176
177
178def copy_location(new_node, old_node):
179 """
180 Copy source location (`lineno`, `col_offset`, `end_lineno`, and `end_col_offset`
181 attributes) from *old_node* to *new_node* if possible, and return *new_node*.
182 """
183 for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
Haibo Huang5eba2b42021-01-22 11:22:02 -0800184 if attr in old_node._attributes and attr in new_node._attributes:
185 value = getattr(old_node, attr, None)
186 # end_lineno and end_col_offset are optional attributes, and they
187 # should be copied whether the value is None or not.
188 if value is not None or (
189 hasattr(old_node, attr) and attr.startswith("end_")
190 ):
191 setattr(new_node, attr, value)
Haibo Huangd8830302020-03-03 10:09:46 -0800192 return new_node
193
194
195def fix_missing_locations(node):
196 """
197 When you compile a node tree with compile(), the compiler expects lineno and
198 col_offset attributes for every node that supports them. This is rather
199 tedious to fill in for generated nodes, so this helper adds these attributes
200 recursively where not already set, by setting them to the values of the
201 parent node. It works recursively starting at *node*.
202 """
203 def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
204 if 'lineno' in node._attributes:
205 if not hasattr(node, 'lineno'):
206 node.lineno = lineno
207 else:
208 lineno = node.lineno
209 if 'end_lineno' in node._attributes:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800210 if getattr(node, 'end_lineno', None) is None:
Haibo Huangd8830302020-03-03 10:09:46 -0800211 node.end_lineno = end_lineno
212 else:
213 end_lineno = node.end_lineno
214 if 'col_offset' in node._attributes:
215 if not hasattr(node, 'col_offset'):
216 node.col_offset = col_offset
217 else:
218 col_offset = node.col_offset
219 if 'end_col_offset' in node._attributes:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800220 if getattr(node, 'end_col_offset', None) is None:
Haibo Huangd8830302020-03-03 10:09:46 -0800221 node.end_col_offset = end_col_offset
222 else:
223 end_col_offset = node.end_col_offset
224 for child in iter_child_nodes(node):
225 _fix(child, lineno, col_offset, end_lineno, end_col_offset)
226 _fix(node, 1, 0, 1, 0)
227 return node
228
229
230def increment_lineno(node, n=1):
231 """
232 Increment the line number and end line number of each node in the tree
233 starting at *node* by *n*. This is useful to "move code" to a different
234 location in a file.
235 """
236 for child in walk(node):
237 if 'lineno' in child._attributes:
238 child.lineno = getattr(child, 'lineno', 0) + n
Haibo Huang5eba2b42021-01-22 11:22:02 -0800239 if (
240 "end_lineno" in child._attributes
241 and (end_lineno := getattr(child, "end_lineno", 0)) is not None
242 ):
243 child.end_lineno = end_lineno + n
Haibo Huangd8830302020-03-03 10:09:46 -0800244 return node
245
246
247def iter_fields(node):
248 """
249 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
250 that is present on *node*.
251 """
252 for field in node._fields:
253 try:
254 yield field, getattr(node, field)
255 except AttributeError:
256 pass
257
258
259def iter_child_nodes(node):
260 """
261 Yield all direct child nodes of *node*, that is, all fields that are nodes
262 and all items of fields that are lists of nodes.
263 """
264 for name, field in iter_fields(node):
265 if isinstance(field, AST):
266 yield field
267 elif isinstance(field, list):
268 for item in field:
269 if isinstance(item, AST):
270 yield item
271
272
273def get_docstring(node, clean=True):
274 """
275 Return the docstring for the given node or None if no docstring can
276 be found. If the node provided does not have docstrings a TypeError
277 will be raised.
278
279 If *clean* is `True`, all tabs are expanded to spaces and any whitespace
280 that can be uniformly removed from the second line onwards is removed.
281 """
282 if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
283 raise TypeError("%r can't have docstrings" % node.__class__.__name__)
284 if not(node.body and isinstance(node.body[0], Expr)):
285 return None
286 node = node.body[0].value
287 if isinstance(node, Str):
288 text = node.s
289 elif isinstance(node, Constant) and isinstance(node.value, str):
290 text = node.value
291 else:
292 return None
293 if clean:
294 import inspect
295 text = inspect.cleandoc(text)
296 return text
297
298
299def _splitlines_no_ff(source):
300 """Split a string into lines ignoring form feed and other chars.
301
302 This mimics how the Python parser splits source code.
303 """
304 idx = 0
305 lines = []
306 next_line = ''
307 while idx < len(source):
308 c = source[idx]
309 next_line += c
310 idx += 1
311 # Keep \r\n together
312 if c == '\r' and idx < len(source) and source[idx] == '\n':
313 next_line += '\n'
314 idx += 1
315 if c in '\r\n':
316 lines.append(next_line)
317 next_line = ''
318
319 if next_line:
320 lines.append(next_line)
321 return lines
322
323
324def _pad_whitespace(source):
Haibo Huang5eba2b42021-01-22 11:22:02 -0800325 r"""Replace all chars except '\f\t' in a line with spaces."""
Haibo Huangd8830302020-03-03 10:09:46 -0800326 result = ''
327 for c in source:
328 if c in '\f\t':
329 result += c
330 else:
331 result += ' '
332 return result
333
334
335def get_source_segment(source, node, *, padded=False):
336 """Get source code segment of the *source* that generated *node*.
337
338 If some location information (`lineno`, `end_lineno`, `col_offset`,
339 or `end_col_offset`) is missing, return None.
340
341 If *padded* is `True`, the first line of a multi-line statement will
342 be padded with spaces to match its original position.
343 """
344 try:
Haibo Huang5eba2b42021-01-22 11:22:02 -0800345 if node.end_lineno is None or node.end_col_offset is None:
346 return None
Haibo Huangd8830302020-03-03 10:09:46 -0800347 lineno = node.lineno - 1
348 end_lineno = node.end_lineno - 1
349 col_offset = node.col_offset
350 end_col_offset = node.end_col_offset
351 except AttributeError:
352 return None
353
354 lines = _splitlines_no_ff(source)
355 if end_lineno == lineno:
356 return lines[lineno].encode()[col_offset:end_col_offset].decode()
357
358 if padded:
359 padding = _pad_whitespace(lines[lineno].encode()[:col_offset].decode())
360 else:
361 padding = ''
362
363 first = padding + lines[lineno].encode()[col_offset:].decode()
364 last = lines[end_lineno].encode()[:end_col_offset].decode()
365 lines = lines[lineno+1:end_lineno]
366
367 lines.insert(0, first)
368 lines.append(last)
369 return ''.join(lines)
370
371
372def walk(node):
373 """
374 Recursively yield all descendant nodes in the tree starting at *node*
375 (including *node* itself), in no specified order. This is useful if you
376 only want to modify nodes in place and don't care about the context.
377 """
378 from collections import deque
379 todo = deque([node])
380 while todo:
381 node = todo.popleft()
382 todo.extend(iter_child_nodes(node))
383 yield node
384
385
386class NodeVisitor(object):
387 """
388 A node visitor base class that walks the abstract syntax tree and calls a
389 visitor function for every node found. This function may return a value
390 which is forwarded by the `visit` method.
391
392 This class is meant to be subclassed, with the subclass adding visitor
393 methods.
394
395 Per default the visitor functions for the nodes are ``'visit_'`` +
396 class name of the node. So a `TryFinally` node visit function would
397 be `visit_TryFinally`. This behavior can be changed by overriding
398 the `visit` method. If no visitor function exists for a node
399 (return value `None`) the `generic_visit` visitor is used instead.
400
401 Don't use the `NodeVisitor` if you want to apply changes to nodes during
402 traversing. For this a special visitor exists (`NodeTransformer`) that
403 allows modifications.
404 """
405
406 def visit(self, node):
407 """Visit a node."""
408 method = 'visit_' + node.__class__.__name__
409 visitor = getattr(self, method, self.generic_visit)
410 return visitor(node)
411
412 def generic_visit(self, node):
413 """Called if no explicit visitor function exists for a node."""
414 for field, value in iter_fields(node):
415 if isinstance(value, list):
416 for item in value:
417 if isinstance(item, AST):
418 self.visit(item)
419 elif isinstance(value, AST):
420 self.visit(value)
421
422 def visit_Constant(self, node):
423 value = node.value
424 type_name = _const_node_type_names.get(type(value))
425 if type_name is None:
426 for cls, name in _const_node_type_names.items():
427 if isinstance(value, cls):
428 type_name = name
429 break
430 if type_name is not None:
431 method = 'visit_' + type_name
432 try:
433 visitor = getattr(self, method)
434 except AttributeError:
435 pass
436 else:
437 import warnings
438 warnings.warn(f"{method} is deprecated; add visit_Constant",
Haibo Huang5eba2b42021-01-22 11:22:02 -0800439 DeprecationWarning, 2)
Haibo Huangd8830302020-03-03 10:09:46 -0800440 return visitor(node)
441 return self.generic_visit(node)
442
443
444class NodeTransformer(NodeVisitor):
445 """
446 A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
447 allows modification of nodes.
448
449 The `NodeTransformer` will walk the AST and use the return value of the
450 visitor methods to replace or remove the old node. If the return value of
451 the visitor method is ``None``, the node will be removed from its location,
452 otherwise it is replaced with the return value. The return value may be the
453 original node in which case no replacement takes place.
454
455 Here is an example transformer that rewrites all occurrences of name lookups
456 (``foo``) to ``data['foo']``::
457
458 class RewriteName(NodeTransformer):
459
460 def visit_Name(self, node):
Haibo Huangf5f93a72020-10-19 15:43:42 -0700461 return Subscript(
Haibo Huangd8830302020-03-03 10:09:46 -0800462 value=Name(id='data', ctx=Load()),
Haibo Huang5eba2b42021-01-22 11:22:02 -0800463 slice=Constant(value=node.id),
Haibo Huangd8830302020-03-03 10:09:46 -0800464 ctx=node.ctx
Haibo Huangf5f93a72020-10-19 15:43:42 -0700465 )
Haibo Huangd8830302020-03-03 10:09:46 -0800466
467 Keep in mind that if the node you're operating on has child nodes you must
468 either transform the child nodes yourself or call the :meth:`generic_visit`
469 method for the node first.
470
471 For nodes that were part of a collection of statements (that applies to all
472 statement nodes), the visitor may also return a list of nodes rather than
473 just a single node.
474
475 Usually you use the transformer like this::
476
477 node = YourTransformer().visit(node)
478 """
479
480 def generic_visit(self, node):
481 for field, old_value in iter_fields(node):
482 if isinstance(old_value, list):
483 new_values = []
484 for value in old_value:
485 if isinstance(value, AST):
486 value = self.visit(value)
487 if value is None:
488 continue
489 elif not isinstance(value, AST):
490 new_values.extend(value)
491 continue
492 new_values.append(value)
493 old_value[:] = new_values
494 elif isinstance(old_value, AST):
495 new_node = self.visit(old_value)
496 if new_node is None:
497 delattr(node, field)
498 else:
499 setattr(node, field, new_node)
500 return node
501
502
Haibo Huang5eba2b42021-01-22 11:22:02 -0800503# If the ast module is loaded more than once, only add deprecated methods once
504if not hasattr(Constant, 'n'):
505 # The following code is for backward compatibility.
506 # It will be removed in future.
Haibo Huangd8830302020-03-03 10:09:46 -0800507
Haibo Huang5eba2b42021-01-22 11:22:02 -0800508 def _getter(self):
509 """Deprecated. Use value instead."""
510 return self.value
Haibo Huangd8830302020-03-03 10:09:46 -0800511
Haibo Huang5eba2b42021-01-22 11:22:02 -0800512 def _setter(self, value):
513 self.value = value
Haibo Huangd8830302020-03-03 10:09:46 -0800514
Haibo Huang5eba2b42021-01-22 11:22:02 -0800515 Constant.n = property(_getter, _setter)
516 Constant.s = property(_getter, _setter)
Haibo Huangd8830302020-03-03 10:09:46 -0800517
518class _ABC(type):
519
Haibo Huang5eba2b42021-01-22 11:22:02 -0800520 def __init__(cls, *args):
521 cls.__doc__ = """Deprecated AST node class. Use ast.Constant instead"""
522
Haibo Huangd8830302020-03-03 10:09:46 -0800523 def __instancecheck__(cls, inst):
524 if not isinstance(inst, Constant):
525 return False
526 if cls in _const_types:
527 try:
528 value = inst.value
529 except AttributeError:
530 return False
531 else:
532 return (
533 isinstance(value, _const_types[cls]) and
534 not isinstance(value, _const_types_not.get(cls, ()))
535 )
536 return type.__instancecheck__(cls, inst)
537
538def _new(cls, *args, **kwargs):
Haibo Huangf5f93a72020-10-19 15:43:42 -0700539 for key in kwargs:
540 if key not in cls._fields:
541 # arbitrary keyword arguments are accepted
542 continue
543 pos = cls._fields.index(key)
544 if pos < len(args):
545 raise TypeError(f"{cls.__name__} got multiple values for argument {key!r}")
Haibo Huangd8830302020-03-03 10:09:46 -0800546 if cls in _const_types:
547 return Constant(*args, **kwargs)
548 return Constant.__new__(cls, *args, **kwargs)
549
550class Num(Constant, metaclass=_ABC):
551 _fields = ('n',)
552 __new__ = _new
553
554class Str(Constant, metaclass=_ABC):
555 _fields = ('s',)
556 __new__ = _new
557
558class Bytes(Constant, metaclass=_ABC):
559 _fields = ('s',)
560 __new__ = _new
561
562class NameConstant(Constant, metaclass=_ABC):
563 __new__ = _new
564
565class Ellipsis(Constant, metaclass=_ABC):
566 _fields = ()
567
568 def __new__(cls, *args, **kwargs):
569 if cls is Ellipsis:
570 return Constant(..., *args, **kwargs)
571 return Constant.__new__(cls, *args, **kwargs)
572
573_const_types = {
574 Num: (int, float, complex),
575 Str: (str,),
576 Bytes: (bytes,),
577 NameConstant: (type(None), bool),
578 Ellipsis: (type(...),),
579}
580_const_types_not = {
581 Num: (bool,),
582}
Haibo Huang5eba2b42021-01-22 11:22:02 -0800583
Haibo Huangd8830302020-03-03 10:09:46 -0800584_const_node_type_names = {
585 bool: 'NameConstant', # should be before int
586 type(None): 'NameConstant',
587 int: 'Num',
588 float: 'Num',
589 complex: 'Num',
590 str: 'Str',
591 bytes: 'Bytes',
592 type(...): 'Ellipsis',
593}
Haibo Huang5eba2b42021-01-22 11:22:02 -0800594
595class slice(AST):
596 """Deprecated AST node class."""
597
598class Index(slice):
599 """Deprecated AST node class. Use the index value directly instead."""
600 def __new__(cls, value, **kwargs):
601 return value
602
603class ExtSlice(slice):
604 """Deprecated AST node class. Use ast.Tuple instead."""
605 def __new__(cls, dims=(), **kwargs):
606 return Tuple(list(dims), Load(), **kwargs)
607
608# If the ast module is loaded more than once, only add deprecated methods once
609if not hasattr(Tuple, 'dims'):
610 # The following code is for backward compatibility.
611 # It will be removed in future.
612
613 def _dims_getter(self):
614 """Deprecated. Use elts instead."""
615 return self.elts
616
617 def _dims_setter(self, value):
618 self.elts = value
619
620 Tuple.dims = property(_dims_getter, _dims_setter)
621
622class Suite(mod):
623 """Deprecated AST node class. Unused in Python 3."""
624
625class AugLoad(expr_context):
626 """Deprecated AST node class. Unused in Python 3."""
627
628class AugStore(expr_context):
629 """Deprecated AST node class. Unused in Python 3."""
630
631class Param(expr_context):
632 """Deprecated AST node class. Unused in Python 3."""
633
634
635# Large float and imaginary literals get turned into infinities in the AST.
636# We unparse those infinities to INFSTR.
637_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
638
639class _Precedence(IntEnum):
640 """Precedence table that originated from python grammar."""
641
642 TUPLE = auto()
643 YIELD = auto() # 'yield', 'yield from'
644 TEST = auto() # 'if'-'else', 'lambda'
645 OR = auto() # 'or'
646 AND = auto() # 'and'
647 NOT = auto() # 'not'
648 CMP = auto() # '<', '>', '==', '>=', '<=', '!=',
649 # 'in', 'not in', 'is', 'is not'
650 EXPR = auto()
651 BOR = EXPR # '|'
652 BXOR = auto() # '^'
653 BAND = auto() # '&'
654 SHIFT = auto() # '<<', '>>'
655 ARITH = auto() # '+', '-'
656 TERM = auto() # '*', '@', '/', '%', '//'
657 FACTOR = auto() # unary '+', '-', '~'
658 POWER = auto() # '**'
659 AWAIT = auto() # 'await'
660 ATOM = auto()
661
662 def next(self):
663 try:
664 return self.__class__(self + 1)
665 except ValueError:
666 return self
667
668
669_SINGLE_QUOTES = ("'", '"')
670_MULTI_QUOTES = ('"""', "'''")
671_ALL_QUOTES = (*_SINGLE_QUOTES, *_MULTI_QUOTES)
672
673class _Unparser(NodeVisitor):
674 """Methods in this class recursively traverse an AST and
675 output source code for the abstract syntax; original formatting
676 is disregarded."""
677
678 def __init__(self, *, _avoid_backslashes=False):
679 self._source = []
680 self._buffer = []
681 self._precedences = {}
682 self._type_ignores = {}
683 self._indent = 0
684 self._avoid_backslashes = _avoid_backslashes
685
686 def interleave(self, inter, f, seq):
687 """Call f on each item in seq, calling inter() in between."""
688 seq = iter(seq)
689 try:
690 f(next(seq))
691 except StopIteration:
692 pass
693 else:
694 for x in seq:
695 inter()
696 f(x)
697
698 def items_view(self, traverser, items):
699 """Traverse and separate the given *items* with a comma and append it to
700 the buffer. If *items* is a single item sequence, a trailing comma
701 will be added."""
702 if len(items) == 1:
703 traverser(items[0])
704 self.write(",")
705 else:
706 self.interleave(lambda: self.write(", "), traverser, items)
707
708 def maybe_newline(self):
709 """Adds a newline if it isn't the start of generated source"""
710 if self._source:
711 self.write("\n")
712
713 def fill(self, text=""):
714 """Indent a piece of text and append it, according to the current
715 indentation level"""
716 self.maybe_newline()
717 self.write(" " * self._indent + text)
718
719 def write(self, text):
720 """Append a piece of text"""
721 self._source.append(text)
722
723 def buffer_writer(self, text):
724 self._buffer.append(text)
725
726 @property
727 def buffer(self):
728 value = "".join(self._buffer)
729 self._buffer.clear()
730 return value
731
732 @contextmanager
733 def block(self, *, extra = None):
734 """A context manager for preparing the source for blocks. It adds
735 the character':', increases the indentation on enter and decreases
736 the indentation on exit. If *extra* is given, it will be directly
737 appended after the colon character.
738 """
739 self.write(":")
740 if extra:
741 self.write(extra)
742 self._indent += 1
743 yield
744 self._indent -= 1
745
746 @contextmanager
747 def delimit(self, start, end):
748 """A context manager for preparing the source for expressions. It adds
749 *start* to the buffer and enters, after exit it adds *end*."""
750
751 self.write(start)
752 yield
753 self.write(end)
754
755 def delimit_if(self, start, end, condition):
756 if condition:
757 return self.delimit(start, end)
758 else:
759 return nullcontext()
760
761 def require_parens(self, precedence, node):
762 """Shortcut to adding precedence related parens"""
763 return self.delimit_if("(", ")", self.get_precedence(node) > precedence)
764
765 def get_precedence(self, node):
766 return self._precedences.get(node, _Precedence.TEST)
767
768 def set_precedence(self, precedence, *nodes):
769 for node in nodes:
770 self._precedences[node] = precedence
771
772 def get_raw_docstring(self, node):
773 """If a docstring node is found in the body of the *node* parameter,
774 return that docstring node, None otherwise.
775
776 Logic mirrored from ``_PyAST_GetDocString``."""
777 if not isinstance(
778 node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
779 ) or len(node.body) < 1:
780 return None
781 node = node.body[0]
782 if not isinstance(node, Expr):
783 return None
784 node = node.value
785 if isinstance(node, Constant) and isinstance(node.value, str):
786 return node
787
788 def get_type_comment(self, node):
789 comment = self._type_ignores.get(node.lineno) or node.type_comment
790 if comment is not None:
791 return f" # type: {comment}"
792
793 def traverse(self, node):
794 if isinstance(node, list):
795 for item in node:
796 self.traverse(item)
797 else:
798 super().visit(node)
799
Yi Kong71199322022-08-30 15:53:45 +0800800 # Note: as visit() resets the output text, do NOT rely on
801 # NodeVisitor.generic_visit to handle any nodes (as it calls back in to
802 # the subclass visit() method, which resets self._source to an empty list)
Haibo Huang5eba2b42021-01-22 11:22:02 -0800803 def visit(self, node):
804 """Outputs a source code string that, if converted back to an ast
805 (using ast.parse) will generate an AST equivalent to *node*"""
806 self._source = []
807 self.traverse(node)
808 return "".join(self._source)
809
810 def _write_docstring_and_traverse_body(self, node):
811 if (docstring := self.get_raw_docstring(node)):
812 self._write_docstring(docstring)
813 self.traverse(node.body[1:])
814 else:
815 self.traverse(node.body)
816
817 def visit_Module(self, node):
818 self._type_ignores = {
819 ignore.lineno: f"ignore{ignore.tag}"
820 for ignore in node.type_ignores
821 }
822 self._write_docstring_and_traverse_body(node)
823 self._type_ignores.clear()
824
825 def visit_FunctionType(self, node):
826 with self.delimit("(", ")"):
827 self.interleave(
828 lambda: self.write(", "), self.traverse, node.argtypes
829 )
830
831 self.write(" -> ")
832 self.traverse(node.returns)
833
834 def visit_Expr(self, node):
835 self.fill()
836 self.set_precedence(_Precedence.YIELD, node.value)
837 self.traverse(node.value)
838
839 def visit_NamedExpr(self, node):
840 with self.require_parens(_Precedence.TUPLE, node):
841 self.set_precedence(_Precedence.ATOM, node.target, node.value)
842 self.traverse(node.target)
843 self.write(" := ")
844 self.traverse(node.value)
845
846 def visit_Import(self, node):
847 self.fill("import ")
848 self.interleave(lambda: self.write(", "), self.traverse, node.names)
849
850 def visit_ImportFrom(self, node):
851 self.fill("from ")
852 self.write("." * node.level)
853 if node.module:
854 self.write(node.module)
855 self.write(" import ")
856 self.interleave(lambda: self.write(", "), self.traverse, node.names)
857
858 def visit_Assign(self, node):
859 self.fill()
860 for target in node.targets:
861 self.traverse(target)
862 self.write(" = ")
863 self.traverse(node.value)
864 if type_comment := self.get_type_comment(node):
865 self.write(type_comment)
866
867 def visit_AugAssign(self, node):
868 self.fill()
869 self.traverse(node.target)
870 self.write(" " + self.binop[node.op.__class__.__name__] + "= ")
871 self.traverse(node.value)
872
873 def visit_AnnAssign(self, node):
874 self.fill()
875 with self.delimit_if("(", ")", not node.simple and isinstance(node.target, Name)):
876 self.traverse(node.target)
877 self.write(": ")
878 self.traverse(node.annotation)
879 if node.value:
880 self.write(" = ")
881 self.traverse(node.value)
882
883 def visit_Return(self, node):
884 self.fill("return")
885 if node.value:
886 self.write(" ")
887 self.traverse(node.value)
888
889 def visit_Pass(self, node):
890 self.fill("pass")
891
892 def visit_Break(self, node):
893 self.fill("break")
894
895 def visit_Continue(self, node):
896 self.fill("continue")
897
898 def visit_Delete(self, node):
899 self.fill("del ")
900 self.interleave(lambda: self.write(", "), self.traverse, node.targets)
901
902 def visit_Assert(self, node):
903 self.fill("assert ")
904 self.traverse(node.test)
905 if node.msg:
906 self.write(", ")
907 self.traverse(node.msg)
908
909 def visit_Global(self, node):
910 self.fill("global ")
911 self.interleave(lambda: self.write(", "), self.write, node.names)
912
913 def visit_Nonlocal(self, node):
914 self.fill("nonlocal ")
915 self.interleave(lambda: self.write(", "), self.write, node.names)
916
917 def visit_Await(self, node):
918 with self.require_parens(_Precedence.AWAIT, node):
919 self.write("await")
920 if node.value:
921 self.write(" ")
922 self.set_precedence(_Precedence.ATOM, node.value)
923 self.traverse(node.value)
924
925 def visit_Yield(self, node):
926 with self.require_parens(_Precedence.YIELD, node):
927 self.write("yield")
928 if node.value:
929 self.write(" ")
930 self.set_precedence(_Precedence.ATOM, node.value)
931 self.traverse(node.value)
932
933 def visit_YieldFrom(self, node):
934 with self.require_parens(_Precedence.YIELD, node):
935 self.write("yield from ")
936 if not node.value:
937 raise ValueError("Node can't be used without a value attribute.")
938 self.set_precedence(_Precedence.ATOM, node.value)
939 self.traverse(node.value)
940
941 def visit_Raise(self, node):
942 self.fill("raise")
943 if not node.exc:
944 if node.cause:
945 raise ValueError(f"Node can't use cause without an exception.")
946 return
947 self.write(" ")
948 self.traverse(node.exc)
949 if node.cause:
950 self.write(" from ")
951 self.traverse(node.cause)
952
953 def visit_Try(self, node):
954 self.fill("try")
955 with self.block():
956 self.traverse(node.body)
957 for ex in node.handlers:
958 self.traverse(ex)
959 if node.orelse:
960 self.fill("else")
961 with self.block():
962 self.traverse(node.orelse)
963 if node.finalbody:
964 self.fill("finally")
965 with self.block():
966 self.traverse(node.finalbody)
967
968 def visit_ExceptHandler(self, node):
969 self.fill("except")
970 if node.type:
971 self.write(" ")
972 self.traverse(node.type)
973 if node.name:
974 self.write(" as ")
975 self.write(node.name)
976 with self.block():
977 self.traverse(node.body)
978
979 def visit_ClassDef(self, node):
980 self.maybe_newline()
981 for deco in node.decorator_list:
982 self.fill("@")
983 self.traverse(deco)
984 self.fill("class " + node.name)
985 with self.delimit_if("(", ")", condition = node.bases or node.keywords):
986 comma = False
987 for e in node.bases:
988 if comma:
989 self.write(", ")
990 else:
991 comma = True
992 self.traverse(e)
993 for e in node.keywords:
994 if comma:
995 self.write(", ")
996 else:
997 comma = True
998 self.traverse(e)
999
1000 with self.block():
1001 self._write_docstring_and_traverse_body(node)
1002
1003 def visit_FunctionDef(self, node):
1004 self._function_helper(node, "def")
1005
1006 def visit_AsyncFunctionDef(self, node):
1007 self._function_helper(node, "async def")
1008
1009 def _function_helper(self, node, fill_suffix):
1010 self.maybe_newline()
1011 for deco in node.decorator_list:
1012 self.fill("@")
1013 self.traverse(deco)
1014 def_str = fill_suffix + " " + node.name
1015 self.fill(def_str)
1016 with self.delimit("(", ")"):
1017 self.traverse(node.args)
1018 if node.returns:
1019 self.write(" -> ")
1020 self.traverse(node.returns)
1021 with self.block(extra=self.get_type_comment(node)):
1022 self._write_docstring_and_traverse_body(node)
1023
1024 def visit_For(self, node):
1025 self._for_helper("for ", node)
1026
1027 def visit_AsyncFor(self, node):
1028 self._for_helper("async for ", node)
1029
1030 def _for_helper(self, fill, node):
1031 self.fill(fill)
1032 self.traverse(node.target)
1033 self.write(" in ")
1034 self.traverse(node.iter)
1035 with self.block(extra=self.get_type_comment(node)):
1036 self.traverse(node.body)
1037 if node.orelse:
1038 self.fill("else")
1039 with self.block():
1040 self.traverse(node.orelse)
1041
1042 def visit_If(self, node):
1043 self.fill("if ")
1044 self.traverse(node.test)
1045 with self.block():
1046 self.traverse(node.body)
1047 # collapse nested ifs into equivalent elifs.
1048 while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], If):
1049 node = node.orelse[0]
1050 self.fill("elif ")
1051 self.traverse(node.test)
1052 with self.block():
1053 self.traverse(node.body)
1054 # final else
1055 if node.orelse:
1056 self.fill("else")
1057 with self.block():
1058 self.traverse(node.orelse)
1059
1060 def visit_While(self, node):
1061 self.fill("while ")
1062 self.traverse(node.test)
1063 with self.block():
1064 self.traverse(node.body)
1065 if node.orelse:
1066 self.fill("else")
1067 with self.block():
1068 self.traverse(node.orelse)
1069
1070 def visit_With(self, node):
1071 self.fill("with ")
1072 self.interleave(lambda: self.write(", "), self.traverse, node.items)
1073 with self.block(extra=self.get_type_comment(node)):
1074 self.traverse(node.body)
1075
1076 def visit_AsyncWith(self, node):
1077 self.fill("async with ")
1078 self.interleave(lambda: self.write(", "), self.traverse, node.items)
1079 with self.block(extra=self.get_type_comment(node)):
1080 self.traverse(node.body)
1081
1082 def _str_literal_helper(
1083 self, string, *, quote_types=_ALL_QUOTES, escape_special_whitespace=False
1084 ):
1085 """Helper for writing string literals, minimizing escapes.
1086 Returns the tuple (string literal to write, possible quote types).
1087 """
1088 def escape_char(c):
1089 # \n and \t are non-printable, but we only escape them if
1090 # escape_special_whitespace is True
1091 if not escape_special_whitespace and c in "\n\t":
1092 return c
1093 # Always escape backslashes and other non-printable characters
1094 if c == "\\" or not c.isprintable():
1095 return c.encode("unicode_escape").decode("ascii")
1096 return c
1097
1098 escaped_string = "".join(map(escape_char, string))
1099 possible_quotes = quote_types
1100 if "\n" in escaped_string:
1101 possible_quotes = [q for q in possible_quotes if q in _MULTI_QUOTES]
1102 possible_quotes = [q for q in possible_quotes if q not in escaped_string]
1103 if not possible_quotes:
1104 # If there aren't any possible_quotes, fallback to using repr
1105 # on the original string. Try to use a quote from quote_types,
1106 # e.g., so that we use triple quotes for docstrings.
1107 string = repr(string)
1108 quote = next((q for q in quote_types if string[0] in q), string[0])
1109 return string[1:-1], [quote]
1110 if escaped_string:
1111 # Sort so that we prefer '''"''' over """\""""
1112 possible_quotes.sort(key=lambda q: q[0] == escaped_string[-1])
1113 # If we're using triple quotes and we'd need to escape a final
1114 # quote, escape it
1115 if possible_quotes[0][0] == escaped_string[-1]:
1116 assert len(possible_quotes[0]) == 3
1117 escaped_string = escaped_string[:-1] + "\\" + escaped_string[-1]
1118 return escaped_string, possible_quotes
1119
1120 def _write_str_avoiding_backslashes(self, string, *, quote_types=_ALL_QUOTES):
1121 """Write string literal value with a best effort attempt to avoid backslashes."""
1122 string, quote_types = self._str_literal_helper(string, quote_types=quote_types)
1123 quote_type = quote_types[0]
1124 self.write(f"{quote_type}{string}{quote_type}")
1125
1126 def visit_JoinedStr(self, node):
1127 self.write("f")
1128 if self._avoid_backslashes:
1129 self._fstring_JoinedStr(node, self.buffer_writer)
1130 self._write_str_avoiding_backslashes(self.buffer)
1131 return
1132
1133 # If we don't need to avoid backslashes globally (i.e., we only need
1134 # to avoid them inside FormattedValues), it's cosmetically preferred
1135 # to use escaped whitespace. That is, it's preferred to use backslashes
1136 # for cases like: f"{x}\n". To accomplish this, we keep track of what
1137 # in our buffer corresponds to FormattedValues and what corresponds to
1138 # Constant parts of the f-string, and allow escapes accordingly.
1139 buffer = []
1140 for value in node.values:
1141 meth = getattr(self, "_fstring_" + type(value).__name__)
1142 meth(value, self.buffer_writer)
1143 buffer.append((self.buffer, isinstance(value, Constant)))
1144 new_buffer = []
1145 quote_types = _ALL_QUOTES
1146 for value, is_constant in buffer:
1147 # Repeatedly narrow down the list of possible quote_types
1148 value, quote_types = self._str_literal_helper(
1149 value, quote_types=quote_types,
1150 escape_special_whitespace=is_constant
1151 )
1152 new_buffer.append(value)
1153 value = "".join(new_buffer)
1154 quote_type = quote_types[0]
1155 self.write(f"{quote_type}{value}{quote_type}")
1156
1157 def visit_FormattedValue(self, node):
1158 self.write("f")
1159 self._fstring_FormattedValue(node, self.buffer_writer)
1160 self._write_str_avoiding_backslashes(self.buffer)
1161
1162 def _fstring_JoinedStr(self, node, write):
1163 for value in node.values:
1164 meth = getattr(self, "_fstring_" + type(value).__name__)
1165 meth(value, write)
1166
1167 def _fstring_Constant(self, node, write):
1168 if not isinstance(node.value, str):
1169 raise ValueError("Constants inside JoinedStr should be a string.")
1170 value = node.value.replace("{", "{{").replace("}", "}}")
1171 write(value)
1172
1173 def _fstring_FormattedValue(self, node, write):
1174 write("{")
1175 unparser = type(self)(_avoid_backslashes=True)
1176 unparser.set_precedence(_Precedence.TEST.next(), node.value)
1177 expr = unparser.visit(node.value)
1178 if expr.startswith("{"):
1179 write(" ") # Separate pair of opening brackets as "{ {"
1180 if "\\" in expr:
1181 raise ValueError("Unable to avoid backslash in f-string expression part")
1182 write(expr)
1183 if node.conversion != -1:
1184 conversion = chr(node.conversion)
1185 if conversion not in "sra":
1186 raise ValueError("Unknown f-string conversion.")
1187 write(f"!{conversion}")
1188 if node.format_spec:
1189 write(":")
1190 meth = getattr(self, "_fstring_" + type(node.format_spec).__name__)
1191 meth(node.format_spec, write)
1192 write("}")
1193
1194 def visit_Name(self, node):
1195 self.write(node.id)
1196
1197 def _write_docstring(self, node):
1198 self.fill()
1199 if node.kind == "u":
1200 self.write("u")
1201 self._write_str_avoiding_backslashes(node.value, quote_types=_MULTI_QUOTES)
1202
1203 def _write_constant(self, value):
1204 if isinstance(value, (float, complex)):
Yi Kong71199322022-08-30 15:53:45 +08001205 # Substitute overflowing decimal literal for AST infinities,
1206 # and inf - inf for NaNs.
1207 self.write(
1208 repr(value)
1209 .replace("inf", _INFSTR)
1210 .replace("nan", f"({_INFSTR}-{_INFSTR})")
1211 )
Haibo Huang5eba2b42021-01-22 11:22:02 -08001212 elif self._avoid_backslashes and isinstance(value, str):
1213 self._write_str_avoiding_backslashes(value)
1214 else:
1215 self.write(repr(value))
1216
1217 def visit_Constant(self, node):
1218 value = node.value
1219 if isinstance(value, tuple):
1220 with self.delimit("(", ")"):
1221 self.items_view(self._write_constant, value)
1222 elif value is ...:
1223 self.write("...")
1224 else:
1225 if node.kind == "u":
1226 self.write("u")
1227 self._write_constant(node.value)
1228
1229 def visit_List(self, node):
1230 with self.delimit("[", "]"):
1231 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
1232
1233 def visit_ListComp(self, node):
1234 with self.delimit("[", "]"):
1235 self.traverse(node.elt)
1236 for gen in node.generators:
1237 self.traverse(gen)
1238
1239 def visit_GeneratorExp(self, node):
1240 with self.delimit("(", ")"):
1241 self.traverse(node.elt)
1242 for gen in node.generators:
1243 self.traverse(gen)
1244
1245 def visit_SetComp(self, node):
1246 with self.delimit("{", "}"):
1247 self.traverse(node.elt)
1248 for gen in node.generators:
1249 self.traverse(gen)
1250
1251 def visit_DictComp(self, node):
1252 with self.delimit("{", "}"):
1253 self.traverse(node.key)
1254 self.write(": ")
1255 self.traverse(node.value)
1256 for gen in node.generators:
1257 self.traverse(gen)
1258
1259 def visit_comprehension(self, node):
1260 if node.is_async:
1261 self.write(" async for ")
1262 else:
1263 self.write(" for ")
1264 self.set_precedence(_Precedence.TUPLE, node.target)
1265 self.traverse(node.target)
1266 self.write(" in ")
1267 self.set_precedence(_Precedence.TEST.next(), node.iter, *node.ifs)
1268 self.traverse(node.iter)
1269 for if_clause in node.ifs:
1270 self.write(" if ")
1271 self.traverse(if_clause)
1272
1273 def visit_IfExp(self, node):
1274 with self.require_parens(_Precedence.TEST, node):
1275 self.set_precedence(_Precedence.TEST.next(), node.body, node.test)
1276 self.traverse(node.body)
1277 self.write(" if ")
1278 self.traverse(node.test)
1279 self.write(" else ")
1280 self.set_precedence(_Precedence.TEST, node.orelse)
1281 self.traverse(node.orelse)
1282
1283 def visit_Set(self, node):
Yi Kong71199322022-08-30 15:53:45 +08001284 if node.elts:
1285 with self.delimit("{", "}"):
1286 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
1287 else:
1288 # `{}` would be interpreted as a dictionary literal, and
1289 # `set` might be shadowed. Thus:
1290 self.write('{*()}')
Haibo Huang5eba2b42021-01-22 11:22:02 -08001291
1292 def visit_Dict(self, node):
1293 def write_key_value_pair(k, v):
1294 self.traverse(k)
1295 self.write(": ")
1296 self.traverse(v)
1297
1298 def write_item(item):
1299 k, v = item
1300 if k is None:
1301 # for dictionary unpacking operator in dicts {**{'y': 2}}
1302 # see PEP 448 for details
1303 self.write("**")
1304 self.set_precedence(_Precedence.EXPR, v)
1305 self.traverse(v)
1306 else:
1307 write_key_value_pair(k, v)
1308
1309 with self.delimit("{", "}"):
1310 self.interleave(
1311 lambda: self.write(", "), write_item, zip(node.keys, node.values)
1312 )
1313
1314 def visit_Tuple(self, node):
1315 with self.delimit("(", ")"):
1316 self.items_view(self.traverse, node.elts)
1317
1318 unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"}
1319 unop_precedence = {
1320 "not": _Precedence.NOT,
1321 "~": _Precedence.FACTOR,
1322 "+": _Precedence.FACTOR,
1323 "-": _Precedence.FACTOR,
1324 }
1325
1326 def visit_UnaryOp(self, node):
1327 operator = self.unop[node.op.__class__.__name__]
1328 operator_precedence = self.unop_precedence[operator]
1329 with self.require_parens(operator_precedence, node):
1330 self.write(operator)
1331 # factor prefixes (+, -, ~) shouldn't be seperated
1332 # from the value they belong, (e.g: +1 instead of + 1)
1333 if operator_precedence is not _Precedence.FACTOR:
1334 self.write(" ")
1335 self.set_precedence(operator_precedence, node.operand)
1336 self.traverse(node.operand)
1337
1338 binop = {
1339 "Add": "+",
1340 "Sub": "-",
1341 "Mult": "*",
1342 "MatMult": "@",
1343 "Div": "/",
1344 "Mod": "%",
1345 "LShift": "<<",
1346 "RShift": ">>",
1347 "BitOr": "|",
1348 "BitXor": "^",
1349 "BitAnd": "&",
1350 "FloorDiv": "//",
1351 "Pow": "**",
1352 }
1353
1354 binop_precedence = {
1355 "+": _Precedence.ARITH,
1356 "-": _Precedence.ARITH,
1357 "*": _Precedence.TERM,
1358 "@": _Precedence.TERM,
1359 "/": _Precedence.TERM,
1360 "%": _Precedence.TERM,
1361 "<<": _Precedence.SHIFT,
1362 ">>": _Precedence.SHIFT,
1363 "|": _Precedence.BOR,
1364 "^": _Precedence.BXOR,
1365 "&": _Precedence.BAND,
1366 "//": _Precedence.TERM,
1367 "**": _Precedence.POWER,
1368 }
1369
1370 binop_rassoc = frozenset(("**",))
1371 def visit_BinOp(self, node):
1372 operator = self.binop[node.op.__class__.__name__]
1373 operator_precedence = self.binop_precedence[operator]
1374 with self.require_parens(operator_precedence, node):
1375 if operator in self.binop_rassoc:
1376 left_precedence = operator_precedence.next()
1377 right_precedence = operator_precedence
1378 else:
1379 left_precedence = operator_precedence
1380 right_precedence = operator_precedence.next()
1381
1382 self.set_precedence(left_precedence, node.left)
1383 self.traverse(node.left)
1384 self.write(f" {operator} ")
1385 self.set_precedence(right_precedence, node.right)
1386 self.traverse(node.right)
1387
1388 cmpops = {
1389 "Eq": "==",
1390 "NotEq": "!=",
1391 "Lt": "<",
1392 "LtE": "<=",
1393 "Gt": ">",
1394 "GtE": ">=",
1395 "Is": "is",
1396 "IsNot": "is not",
1397 "In": "in",
1398 "NotIn": "not in",
1399 }
1400
1401 def visit_Compare(self, node):
1402 with self.require_parens(_Precedence.CMP, node):
1403 self.set_precedence(_Precedence.CMP.next(), node.left, *node.comparators)
1404 self.traverse(node.left)
1405 for o, e in zip(node.ops, node.comparators):
1406 self.write(" " + self.cmpops[o.__class__.__name__] + " ")
1407 self.traverse(e)
1408
1409 boolops = {"And": "and", "Or": "or"}
1410 boolop_precedence = {"and": _Precedence.AND, "or": _Precedence.OR}
1411
1412 def visit_BoolOp(self, node):
1413 operator = self.boolops[node.op.__class__.__name__]
1414 operator_precedence = self.boolop_precedence[operator]
1415
1416 def increasing_level_traverse(node):
1417 nonlocal operator_precedence
1418 operator_precedence = operator_precedence.next()
1419 self.set_precedence(operator_precedence, node)
1420 self.traverse(node)
1421
1422 with self.require_parens(operator_precedence, node):
1423 s = f" {operator} "
1424 self.interleave(lambda: self.write(s), increasing_level_traverse, node.values)
1425
1426 def visit_Attribute(self, node):
1427 self.set_precedence(_Precedence.ATOM, node.value)
1428 self.traverse(node.value)
1429 # Special case: 3.__abs__() is a syntax error, so if node.value
1430 # is an integer literal then we need to either parenthesize
1431 # it or add an extra space to get 3 .__abs__().
1432 if isinstance(node.value, Constant) and isinstance(node.value.value, int):
1433 self.write(" ")
1434 self.write(".")
1435 self.write(node.attr)
1436
1437 def visit_Call(self, node):
1438 self.set_precedence(_Precedence.ATOM, node.func)
1439 self.traverse(node.func)
1440 with self.delimit("(", ")"):
1441 comma = False
1442 for e in node.args:
1443 if comma:
1444 self.write(", ")
1445 else:
1446 comma = True
1447 self.traverse(e)
1448 for e in node.keywords:
1449 if comma:
1450 self.write(", ")
1451 else:
1452 comma = True
1453 self.traverse(e)
1454
1455 def visit_Subscript(self, node):
1456 def is_simple_tuple(slice_value):
Yi Kong71199322022-08-30 15:53:45 +08001457 # when unparsing a non-empty tuple, the parentheses can be safely
Haibo Huang5eba2b42021-01-22 11:22:02 -08001458 # omitted if there aren't any elements that explicitly requires
Yi Kong71199322022-08-30 15:53:45 +08001459 # parentheses (such as starred expressions).
Haibo Huang5eba2b42021-01-22 11:22:02 -08001460 return (
1461 isinstance(slice_value, Tuple)
1462 and slice_value.elts
1463 and not any(isinstance(elt, Starred) for elt in slice_value.elts)
1464 )
1465
1466 self.set_precedence(_Precedence.ATOM, node.value)
1467 self.traverse(node.value)
1468 with self.delimit("[", "]"):
1469 if is_simple_tuple(node.slice):
1470 self.items_view(self.traverse, node.slice.elts)
1471 else:
1472 self.traverse(node.slice)
1473
1474 def visit_Starred(self, node):
1475 self.write("*")
1476 self.set_precedence(_Precedence.EXPR, node.value)
1477 self.traverse(node.value)
1478
1479 def visit_Ellipsis(self, node):
1480 self.write("...")
1481
1482 def visit_Slice(self, node):
1483 if node.lower:
1484 self.traverse(node.lower)
1485 self.write(":")
1486 if node.upper:
1487 self.traverse(node.upper)
1488 if node.step:
1489 self.write(":")
1490 self.traverse(node.step)
1491
Yi Kong71199322022-08-30 15:53:45 +08001492 def visit_Match(self, node):
1493 self.fill("match ")
1494 self.traverse(node.subject)
1495 with self.block():
1496 for case in node.cases:
1497 self.traverse(case)
1498
Haibo Huang5eba2b42021-01-22 11:22:02 -08001499 def visit_arg(self, node):
1500 self.write(node.arg)
1501 if node.annotation:
1502 self.write(": ")
1503 self.traverse(node.annotation)
1504
1505 def visit_arguments(self, node):
1506 first = True
1507 # normal arguments
1508 all_args = node.posonlyargs + node.args
1509 defaults = [None] * (len(all_args) - len(node.defaults)) + node.defaults
1510 for index, elements in enumerate(zip(all_args, defaults), 1):
1511 a, d = elements
1512 if first:
1513 first = False
1514 else:
1515 self.write(", ")
1516 self.traverse(a)
1517 if d:
1518 self.write("=")
1519 self.traverse(d)
1520 if index == len(node.posonlyargs):
1521 self.write(", /")
1522
1523 # varargs, or bare '*' if no varargs but keyword-only arguments present
1524 if node.vararg or node.kwonlyargs:
1525 if first:
1526 first = False
1527 else:
1528 self.write(", ")
1529 self.write("*")
1530 if node.vararg:
1531 self.write(node.vararg.arg)
1532 if node.vararg.annotation:
1533 self.write(": ")
1534 self.traverse(node.vararg.annotation)
1535
1536 # keyword-only arguments
1537 if node.kwonlyargs:
1538 for a, d in zip(node.kwonlyargs, node.kw_defaults):
1539 self.write(", ")
1540 self.traverse(a)
1541 if d:
1542 self.write("=")
1543 self.traverse(d)
1544
1545 # kwargs
1546 if node.kwarg:
1547 if first:
1548 first = False
1549 else:
1550 self.write(", ")
1551 self.write("**" + node.kwarg.arg)
1552 if node.kwarg.annotation:
1553 self.write(": ")
1554 self.traverse(node.kwarg.annotation)
1555
1556 def visit_keyword(self, node):
1557 if node.arg is None:
1558 self.write("**")
1559 else:
1560 self.write(node.arg)
1561 self.write("=")
1562 self.traverse(node.value)
1563
1564 def visit_Lambda(self, node):
1565 with self.require_parens(_Precedence.TEST, node):
1566 self.write("lambda ")
1567 self.traverse(node.args)
1568 self.write(": ")
1569 self.set_precedence(_Precedence.TEST, node.body)
1570 self.traverse(node.body)
1571
1572 def visit_alias(self, node):
1573 self.write(node.name)
1574 if node.asname:
1575 self.write(" as " + node.asname)
1576
1577 def visit_withitem(self, node):
1578 self.traverse(node.context_expr)
1579 if node.optional_vars:
1580 self.write(" as ")
1581 self.traverse(node.optional_vars)
1582
Yi Kong71199322022-08-30 15:53:45 +08001583 def visit_match_case(self, node):
1584 self.fill("case ")
1585 self.traverse(node.pattern)
1586 if node.guard:
1587 self.write(" if ")
1588 self.traverse(node.guard)
1589 with self.block():
1590 self.traverse(node.body)
1591
1592 def visit_MatchValue(self, node):
1593 self.traverse(node.value)
1594
1595 def visit_MatchSingleton(self, node):
1596 self._write_constant(node.value)
1597
1598 def visit_MatchSequence(self, node):
1599 with self.delimit("[", "]"):
1600 self.interleave(
1601 lambda: self.write(", "), self.traverse, node.patterns
1602 )
1603
1604 def visit_MatchStar(self, node):
1605 name = node.name
1606 if name is None:
1607 name = "_"
1608 self.write(f"*{name}")
1609
1610 def visit_MatchMapping(self, node):
1611 def write_key_pattern_pair(pair):
1612 k, p = pair
1613 self.traverse(k)
1614 self.write(": ")
1615 self.traverse(p)
1616
1617 with self.delimit("{", "}"):
1618 keys = node.keys
1619 self.interleave(
1620 lambda: self.write(", "),
1621 write_key_pattern_pair,
1622 zip(keys, node.patterns, strict=True),
1623 )
1624 rest = node.rest
1625 if rest is not None:
1626 if keys:
1627 self.write(", ")
1628 self.write(f"**{rest}")
1629
1630 def visit_MatchClass(self, node):
1631 self.set_precedence(_Precedence.ATOM, node.cls)
1632 self.traverse(node.cls)
1633 with self.delimit("(", ")"):
1634 patterns = node.patterns
1635 self.interleave(
1636 lambda: self.write(", "), self.traverse, patterns
1637 )
1638 attrs = node.kwd_attrs
1639 if attrs:
1640 def write_attr_pattern(pair):
1641 attr, pattern = pair
1642 self.write(f"{attr}=")
1643 self.traverse(pattern)
1644
1645 if patterns:
1646 self.write(", ")
1647 self.interleave(
1648 lambda: self.write(", "),
1649 write_attr_pattern,
1650 zip(attrs, node.kwd_patterns, strict=True),
1651 )
1652
1653 def visit_MatchAs(self, node):
1654 name = node.name
1655 pattern = node.pattern
1656 if name is None:
1657 self.write("_")
1658 elif pattern is None:
1659 self.write(node.name)
1660 else:
1661 with self.require_parens(_Precedence.TEST, node):
1662 self.set_precedence(_Precedence.BOR, node.pattern)
1663 self.traverse(node.pattern)
1664 self.write(f" as {node.name}")
1665
1666 def visit_MatchOr(self, node):
1667 with self.require_parens(_Precedence.BOR, node):
1668 self.set_precedence(_Precedence.BOR.next(), *node.patterns)
1669 self.interleave(lambda: self.write(" | "), self.traverse, node.patterns)
1670
Haibo Huang5eba2b42021-01-22 11:22:02 -08001671def unparse(ast_obj):
1672 unparser = _Unparser()
1673 return unparser.visit(ast_obj)
1674
1675
1676def main():
1677 import argparse
1678
1679 parser = argparse.ArgumentParser(prog='python -m ast')
1680 parser.add_argument('infile', type=argparse.FileType(mode='rb'), nargs='?',
1681 default='-',
1682 help='the file to parse; defaults to stdin')
1683 parser.add_argument('-m', '--mode', default='exec',
1684 choices=('exec', 'single', 'eval', 'func_type'),
1685 help='specify what kind of code must be parsed')
1686 parser.add_argument('--no-type-comments', default=True, action='store_false',
1687 help="don't add information about type comments")
1688 parser.add_argument('-a', '--include-attributes', action='store_true',
1689 help='include attributes such as line numbers and '
1690 'column offsets')
1691 parser.add_argument('-i', '--indent', type=int, default=3,
1692 help='indentation of nodes (number of spaces)')
1693 args = parser.parse_args()
1694
1695 with args.infile as infile:
1696 source = infile.read()
1697 tree = parse(source, args.infile.name, args.mode, type_comments=args.no_type_comments)
1698 print(dump(tree, include_attributes=args.include_attributes, indent=args.indent))
1699
1700if __name__ == '__main__':
1701 main()