blob: 0c53e5c5712f5ed2d8cbd9997ef1cceebe698596 [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
Ethan Furmana02cb472021-04-21 10:20:44 -070030from enum import IntEnum, auto, _simple_enum
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):
Batuhan Taskayae799aa82020-10-04 03:46:44 +030062 node_or_string = parse(node_or_string.lstrip(" \t"), mode='eval')
Georg Brandl0c77a822008-06-10 16:37:50 +000063 if isinstance(node_or_string, Expression):
64 node_or_string = node_or_string.body
Curtis Bucherc21c5122020-05-05 12:40:56 -070065 def _raise_malformed_node(node):
Irit Katriel586f3db2020-12-25 17:04:31 +000066 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}')
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020070 def _convert_num(node):
Curtis Bucherc21c5122020-05-05 12:40:56 -070071 if not isinstance(node, Constant) or type(node.value) not in (int, float, complex):
72 _raise_malformed_node(node)
73 return node.value
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020074 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)
Georg Brandl0c77a822008-06-10 16:37:50 +000082 def _convert(node):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010083 if isinstance(node, Constant):
84 return node.value
Georg Brandl0c77a822008-06-10 16:37:50 +000085 elif isinstance(node, Tuple):
86 return tuple(map(_convert, node.elts))
87 elif isinstance(node, List):
88 return list(map(_convert, node.elts))
Georg Brandl492f3fc2010-07-11 09:41:21 +000089 elif isinstance(node, Set):
90 return set(map(_convert, node.elts))
Raymond Hettinger4fcf5c12020-01-02 22:21:18 -070091 elif (isinstance(node, Call) and isinstance(node.func, Name) and
92 node.func.id == 'set' and node.args == node.keywords == []):
93 return set()
Georg Brandl0c77a822008-06-10 16:37:50 +000094 elif isinstance(node, Dict):
Curtis Bucherc21c5122020-05-05 12:40:56 -070095 if len(node.keys) != len(node.values):
96 _raise_malformed_node(node)
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +020097 return dict(zip(map(_convert, node.keys),
98 map(_convert, node.values)))
Victor Stinnerf2c1aa12016-01-26 00:40:57 +010099 elif isinstance(node, BinOp) and isinstance(node.op, (Add, Sub)):
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +0200100 left = _convert_signed_num(node.left)
101 right = _convert_num(node.right)
102 if isinstance(left, (int, float)) and isinstance(right, complex):
Victor Stinnerf2c1aa12016-01-26 00:40:57 +0100103 if isinstance(node.op, Add):
104 return left + right
105 else:
106 return left - right
Serhiy Storchakad8ac4d12018-01-04 11:15:39 +0200107 return _convert_signed_num(node)
Georg Brandl0c77a822008-06-10 16:37:50 +0000108 return _convert(node_or_string)
109
110
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300111def dump(node, annotate_fields=True, include_attributes=False, *, indent=None):
Georg Brandl0c77a822008-06-10 16:37:50 +0000112 """
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300113 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
Benjamin Petersondcf97b92008-07-02 17:30:14 +0000118 numbers and column offsets are not dumped by default. If this is wanted,
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300119 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.
Georg Brandl0c77a822008-06-10 16:37:50 +0000122 """
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300123 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 = ', '
Georg Brandl0c77a822008-06-10 16:37:50 +0000131 if isinstance(node, AST):
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200132 cls = type(node)
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300133 args = []
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300134 allsimple = True
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300135 keywords = annotate_fields
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200136 for name in node._fields:
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300137 try:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200138 value = getattr(node, name)
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300139 except AttributeError:
140 keywords = True
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200141 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))
Serhiy Storchakae64f9482019-08-29 09:30:23 +0300149 else:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200150 args.append(value)
151 if include_attributes and node._attributes:
152 for name in node._attributes:
153 try:
154 value = getattr(node, name)
155 except AttributeError:
156 continue
157 if value is None and getattr(cls, name, ...) is None:
158 continue
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300159 value, simple = _format(value, level)
160 allsimple = allsimple and simple
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200161 args.append('%s=%s' % (name, value))
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300162 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
Georg Brandl0c77a822008-06-10 16:37:50 +0000165 elif isinstance(node, list):
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300166 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
Georg Brandl0c77a822008-06-10 16:37:50 +0000171 if not isinstance(node, AST):
172 raise TypeError('expected AST, got %r' % node.__class__.__name__)
Serhiy Storchaka850573b2019-09-09 19:33:13 +0300173 if indent is not None and not isinstance(indent, str):
174 indent = ' ' * indent
175 return _format(node)[0]
Georg Brandl0c77a822008-06-10 16:37:50 +0000176
177
178def copy_location(new_node, old_node):
179 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000180 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*.
Georg Brandl0c77a822008-06-10 16:37:50 +0000182 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000183 for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200184 if attr in old_node._attributes and attr in new_node._attributes:
185 value = getattr(old_node, attr, None)
Batuhan Taskaya8f4380d2020-08-05 16:32:32 +0300186 # 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 ):
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200191 setattr(new_node, attr, value)
Georg Brandl0c77a822008-06-10 16:37:50 +0000192 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 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000203 def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
Georg Brandl0c77a822008-06-10 16:37:50 +0000204 if 'lineno' in node._attributes:
205 if not hasattr(node, 'lineno'):
206 node.lineno = lineno
207 else:
208 lineno = node.lineno
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000209 if 'end_lineno' in node._attributes:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200210 if getattr(node, 'end_lineno', None) is None:
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000211 node.end_lineno = end_lineno
212 else:
213 end_lineno = node.end_lineno
Georg Brandl0c77a822008-06-10 16:37:50 +0000214 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
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000219 if 'end_col_offset' in node._attributes:
Serhiy Storchakab7e95252020-03-10 00:07:47 +0200220 if getattr(node, 'end_col_offset', None) is None:
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000221 node.end_col_offset = end_col_offset
222 else:
223 end_col_offset = node.end_col_offset
Georg Brandl0c77a822008-06-10 16:37:50 +0000224 for child in iter_child_nodes(node):
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000225 _fix(child, lineno, col_offset, end_lineno, end_col_offset)
226 _fix(node, 1, 0, 1, 0)
Georg Brandl0c77a822008-06-10 16:37:50 +0000227 return node
228
229
230def increment_lineno(node, n=1):
231 """
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000232 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.
Georg Brandl0c77a822008-06-10 16:37:50 +0000235 """
Georg Brandl0c77a822008-06-10 16:37:50 +0000236 for child in walk(node):
237 if 'lineno' in child._attributes:
238 child.lineno = getattr(child, 'lineno', 0) + n
Batuhan Taskaya8f4380d2020-08-05 16:32:32 +0300239 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
Georg Brandl0c77a822008-06-10 16:37:50 +0000244 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.
Matthias Bussonnier41cea702017-02-23 22:44:19 -0800278
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.
Georg Brandl0c77a822008-06-10 16:37:50 +0000281 """
Yury Selivanov2f07a662015-07-23 08:54:35 +0300282 if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
Georg Brandl0c77a822008-06-10 16:37:50 +0000283 raise TypeError("%r can't have docstrings" % node.__class__.__name__)
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300284 if not(node.body and isinstance(node.body[0], Expr)):
Serhiy Storchaka73cbe7a2018-05-29 12:04:55 +0300285 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
Serhiy Storchaka08f127a2018-06-15 11:05:15 +0300293 if clean:
Victor Stinnerf2c1aa12016-01-26 00:40:57 +0100294 import inspect
295 text = inspect.cleandoc(text)
296 return text
Georg Brandl0c77a822008-06-10 16:37:50 +0000297
298
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000299def _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):
mpheathfbeba8f2020-02-14 04:32:09 +1000325 r"""Replace all chars except '\f\t' in a line with spaces."""
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000326 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:
Irit Katriele6578a22020-05-18 19:14:12 +0100345 if node.end_lineno is None or node.end_col_offset is None:
346 return None
Ivan Levkivskyi9932a222019-01-22 11:18:22 +0000347 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
Georg Brandl0c77a822008-06-10 16:37:50 +0000372def walk(node):
373 """
Georg Brandl619e7ba2011-01-09 07:38:51 +0000374 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.
Georg Brandl0c77a822008-06-10 16:37:50 +0000377 """
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
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300422 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",
439 DeprecationWarning, 2)
440 return visitor(node)
441 return self.generic_visit(node)
442
Georg Brandl0c77a822008-06-10 16:37:50 +0000443
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):
Pablo Galindoc00c86b2020-03-12 00:48:19 +0000461 return Subscript(
Georg Brandl0c77a822008-06-10 16:37:50 +0000462 value=Name(id='data', ctx=Load()),
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200463 slice=Constant(value=node.id),
Georg Brandl0c77a822008-06-10 16:37:50 +0000464 ctx=node.ctx
Pablo Galindoc00c86b2020-03-12 00:48:19 +0000465 )
Georg Brandl0c77a822008-06-10 16:37:50 +0000466
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):
Georg Brandl0c77a822008-06-10 16:37:50 +0000482 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
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300501
502
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200503# 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.
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300507
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200508 def _getter(self):
509 """Deprecated. Use value instead."""
510 return self.value
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300511
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200512 def _setter(self, value):
513 self.value = value
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300514
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200515 Constant.n = property(_getter, _setter)
516 Constant.s = property(_getter, _setter)
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300517
518class _ABC(type):
519
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200520 def __init__(cls, *args):
521 cls.__doc__ = """Deprecated AST node class. Use ast.Constant instead"""
522
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300523 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:
Anthony Sottile74176222019-01-18 11:30:28 -0800532 return (
533 isinstance(value, _const_types[cls]) and
534 not isinstance(value, _const_types_not.get(cls, ()))
535 )
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300536 return type.__instancecheck__(cls, inst)
537
538def _new(cls, *args, **kwargs):
Rémi Lapeyrec73914a2020-05-24 23:12:57 +0200539 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}")
Serhiy Storchaka3f228112018-09-27 17:42:37 +0300546 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}
Anthony Sottile74176222019-01-18 11:30:28 -0800580_const_types_not = {
581 Num: (bool,),
582}
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200583
Serhiy Storchakac3ea41e2019-08-26 10:13:19 +0300584_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}
Serhiy Storchaka832e8642019-09-09 23:36:13 +0300594
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200595class slice(AST):
596 """Deprecated AST node class."""
597
598class Index(slice):
599 """Deprecated AST node class. Use the index value directly instead."""
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200600 def __new__(cls, value, **kwargs):
601 return value
602
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200603class ExtSlice(slice):
604 """Deprecated AST node class. Use ast.Tuple instead."""
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200605 def __new__(cls, dims=(), **kwargs):
606 return Tuple(list(dims), Load(), **kwargs)
607
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200608# 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.
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200612
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200613 def _dims_getter(self):
614 """Deprecated. Use elts instead."""
615 return self.elts
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200616
Victor Stinnere5fbe0c2020-09-15 18:03:34 +0200617 def _dims_setter(self, value):
618 self.elts = value
619
620 Tuple.dims = property(_dims_getter, _dims_setter)
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200621
Serhiy Storchakabace59d2020-03-22 20:33:34 +0200622class 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
Serhiy Storchaka13d52c22020-03-10 18:52:34 +0200634
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000635# 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
Ethan Furmana02cb472021-04-21 10:20:44 -0700639@_simple_enum(IntEnum)
640class _Precedence:
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300641 """Precedence table that originated from python grammar."""
642
643 TUPLE = auto()
644 YIELD = auto() # 'yield', 'yield from'
645 TEST = auto() # 'if'-'else', 'lambda'
646 OR = auto() # 'or'
647 AND = auto() # 'and'
648 NOT = auto() # 'not'
649 CMP = auto() # '<', '>', '==', '>=', '<=', '!=',
650 # 'in', 'not in', 'is', 'is not'
651 EXPR = auto()
652 BOR = EXPR # '|'
653 BXOR = auto() # '^'
654 BAND = auto() # '&'
655 SHIFT = auto() # '<<', '>>'
656 ARITH = auto() # '+', '-'
657 TERM = auto() # '*', '@', '/', '%', '//'
658 FACTOR = auto() # unary '+', '-', '~'
659 POWER = auto() # '**'
660 AWAIT = auto() # 'await'
661 ATOM = auto()
662
663 def next(self):
664 try:
665 return self.__class__(self + 1)
666 except ValueError:
667 return self
668
Shantanua993e902020-11-20 13:16:42 -0800669
670_SINGLE_QUOTES = ("'", '"')
671_MULTI_QUOTES = ('"""', "'''")
672_ALL_QUOTES = (*_SINGLE_QUOTES, *_MULTI_QUOTES)
673
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000674class _Unparser(NodeVisitor):
675 """Methods in this class recursively traverse an AST and
676 output source code for the abstract syntax; original formatting
677 is disregarded."""
678
Shantanua993e902020-11-20 13:16:42 -0800679 def __init__(self, *, _avoid_backslashes=False):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000680 self._source = []
681 self._buffer = []
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300682 self._precedences = {}
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300683 self._type_ignores = {}
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000684 self._indent = 0
Shantanua993e902020-11-20 13:16:42 -0800685 self._avoid_backslashes = _avoid_backslashes
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000686
687 def interleave(self, inter, f, seq):
688 """Call f on each item in seq, calling inter() in between."""
689 seq = iter(seq)
690 try:
691 f(next(seq))
692 except StopIteration:
693 pass
694 else:
695 for x in seq:
696 inter()
697 f(x)
698
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +0300699 def items_view(self, traverser, items):
700 """Traverse and separate the given *items* with a comma and append it to
701 the buffer. If *items* is a single item sequence, a trailing comma
702 will be added."""
703 if len(items) == 1:
704 traverser(items[0])
705 self.write(",")
706 else:
707 self.interleave(lambda: self.write(", "), traverser, items)
708
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300709 def maybe_newline(self):
710 """Adds a newline if it isn't the start of generated source"""
711 if self._source:
712 self.write("\n")
713
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000714 def fill(self, text=""):
715 """Indent a piece of text and append it, according to the current
716 indentation level"""
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300717 self.maybe_newline()
718 self.write(" " * self._indent + text)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000719
720 def write(self, text):
721 """Append a piece of text"""
722 self._source.append(text)
723
724 def buffer_writer(self, text):
725 self._buffer.append(text)
726
727 @property
728 def buffer(self):
729 value = "".join(self._buffer)
730 self._buffer.clear()
731 return value
732
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000733 @contextmanager
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300734 def block(self, *, extra = None):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000735 """A context manager for preparing the source for blocks. It adds
736 the character':', increases the indentation on enter and decreases
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300737 the indentation on exit. If *extra* is given, it will be directly
738 appended after the colon character.
739 """
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000740 self.write(":")
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300741 if extra:
742 self.write(extra)
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000743 self._indent += 1
744 yield
745 self._indent -= 1
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000746
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300747 @contextmanager
748 def delimit(self, start, end):
749 """A context manager for preparing the source for expressions. It adds
750 *start* to the buffer and enters, after exit it adds *end*."""
751
752 self.write(start)
753 yield
754 self.write(end)
755
756 def delimit_if(self, start, end, condition):
757 if condition:
758 return self.delimit(start, end)
759 else:
760 return nullcontext()
761
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300762 def require_parens(self, precedence, node):
763 """Shortcut to adding precedence related parens"""
764 return self.delimit_if("(", ")", self.get_precedence(node) > precedence)
765
766 def get_precedence(self, node):
767 return self._precedences.get(node, _Precedence.TEST)
768
769 def set_precedence(self, precedence, *nodes):
770 for node in nodes:
771 self._precedences[node] = precedence
772
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300773 def get_raw_docstring(self, node):
774 """If a docstring node is found in the body of the *node* parameter,
775 return that docstring node, None otherwise.
776
777 Logic mirrored from ``_PyAST_GetDocString``."""
778 if not isinstance(
779 node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
780 ) or len(node.body) < 1:
781 return None
782 node = node.body[0]
783 if not isinstance(node, Expr):
784 return None
785 node = node.value
786 if isinstance(node, Constant) and isinstance(node.value, str):
787 return node
788
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300789 def get_type_comment(self, node):
790 comment = self._type_ignores.get(node.lineno) or node.type_comment
791 if comment is not None:
792 return f" # type: {comment}"
793
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000794 def traverse(self, node):
795 if isinstance(node, list):
796 for item in node:
797 self.traverse(item)
798 else:
799 super().visit(node)
800
Nick Coghlan1e7b8582021-04-29 15:58:44 +1000801 # Note: as visit() resets the output text, do NOT rely on
802 # NodeVisitor.generic_visit to handle any nodes (as it calls back in to
803 # the subclass visit() method, which resets self._source to an empty list)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000804 def visit(self, node):
805 """Outputs a source code string that, if converted back to an ast
806 (using ast.parse) will generate an AST equivalent to *node*"""
807 self._source = []
808 self.traverse(node)
809 return "".join(self._source)
810
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300811 def _write_docstring_and_traverse_body(self, node):
812 if (docstring := self.get_raw_docstring(node)):
813 self._write_docstring(docstring)
814 self.traverse(node.body[1:])
815 else:
816 self.traverse(node.body)
817
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000818 def visit_Module(self, node):
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300819 self._type_ignores = {
820 ignore.lineno: f"ignore{ignore.tag}"
821 for ignore in node.type_ignores
822 }
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300823 self._write_docstring_and_traverse_body(node)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300824 self._type_ignores.clear()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000825
Batuhan Taşkaya5b66ec12020-03-15 22:56:57 +0300826 def visit_FunctionType(self, node):
827 with self.delimit("(", ")"):
828 self.interleave(
829 lambda: self.write(", "), self.traverse, node.argtypes
830 )
831
832 self.write(" -> ")
833 self.traverse(node.returns)
834
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000835 def visit_Expr(self, node):
836 self.fill()
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300837 self.set_precedence(_Precedence.YIELD, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000838 self.traverse(node.value)
839
840 def visit_NamedExpr(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300841 with self.require_parens(_Precedence.TUPLE, node):
842 self.set_precedence(_Precedence.ATOM, node.target, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300843 self.traverse(node.target)
844 self.write(" := ")
845 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000846
847 def visit_Import(self, node):
848 self.fill("import ")
849 self.interleave(lambda: self.write(", "), self.traverse, node.names)
850
851 def visit_ImportFrom(self, node):
852 self.fill("from ")
853 self.write("." * node.level)
854 if node.module:
855 self.write(node.module)
856 self.write(" import ")
857 self.interleave(lambda: self.write(", "), self.traverse, node.names)
858
859 def visit_Assign(self, node):
860 self.fill()
861 for target in node.targets:
862 self.traverse(target)
863 self.write(" = ")
864 self.traverse(node.value)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300865 if type_comment := self.get_type_comment(node):
866 self.write(type_comment)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000867
868 def visit_AugAssign(self, node):
869 self.fill()
870 self.traverse(node.target)
871 self.write(" " + self.binop[node.op.__class__.__name__] + "= ")
872 self.traverse(node.value)
873
874 def visit_AnnAssign(self, node):
875 self.fill()
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300876 with self.delimit_if("(", ")", not node.simple and isinstance(node.target, Name)):
877 self.traverse(node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000878 self.write(": ")
879 self.traverse(node.annotation)
880 if node.value:
881 self.write(" = ")
882 self.traverse(node.value)
883
884 def visit_Return(self, node):
885 self.fill("return")
886 if node.value:
887 self.write(" ")
888 self.traverse(node.value)
889
890 def visit_Pass(self, node):
891 self.fill("pass")
892
893 def visit_Break(self, node):
894 self.fill("break")
895
896 def visit_Continue(self, node):
897 self.fill("continue")
898
899 def visit_Delete(self, node):
900 self.fill("del ")
901 self.interleave(lambda: self.write(", "), self.traverse, node.targets)
902
903 def visit_Assert(self, node):
904 self.fill("assert ")
905 self.traverse(node.test)
906 if node.msg:
907 self.write(", ")
908 self.traverse(node.msg)
909
910 def visit_Global(self, node):
911 self.fill("global ")
912 self.interleave(lambda: self.write(", "), self.write, node.names)
913
914 def visit_Nonlocal(self, node):
915 self.fill("nonlocal ")
916 self.interleave(lambda: self.write(", "), self.write, node.names)
917
918 def visit_Await(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300919 with self.require_parens(_Precedence.AWAIT, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300920 self.write("await")
921 if node.value:
922 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300923 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300924 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000925
926 def visit_Yield(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300927 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300928 self.write("yield")
929 if node.value:
930 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300931 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300932 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000933
934 def visit_YieldFrom(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300935 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300936 self.write("yield from ")
937 if not node.value:
938 raise ValueError("Node can't be used without a value attribute.")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300939 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300940 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000941
942 def visit_Raise(self, node):
943 self.fill("raise")
944 if not node.exc:
945 if node.cause:
946 raise ValueError(f"Node can't use cause without an exception.")
947 return
948 self.write(" ")
949 self.traverse(node.exc)
950 if node.cause:
951 self.write(" from ")
952 self.traverse(node.cause)
953
954 def visit_Try(self, node):
955 self.fill("try")
956 with self.block():
957 self.traverse(node.body)
958 for ex in node.handlers:
959 self.traverse(ex)
960 if node.orelse:
961 self.fill("else")
962 with self.block():
963 self.traverse(node.orelse)
964 if node.finalbody:
965 self.fill("finally")
966 with self.block():
967 self.traverse(node.finalbody)
968
969 def visit_ExceptHandler(self, node):
970 self.fill("except")
971 if node.type:
972 self.write(" ")
973 self.traverse(node.type)
974 if node.name:
975 self.write(" as ")
976 self.write(node.name)
977 with self.block():
978 self.traverse(node.body)
979
980 def visit_ClassDef(self, node):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300981 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000982 for deco in node.decorator_list:
983 self.fill("@")
984 self.traverse(deco)
985 self.fill("class " + node.name)
Batuhan Taskaya25160cd2020-05-17 00:53:25 +0300986 with self.delimit_if("(", ")", condition = node.bases or node.keywords):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300987 comma = False
988 for e in node.bases:
989 if comma:
990 self.write(", ")
991 else:
992 comma = True
993 self.traverse(e)
994 for e in node.keywords:
995 if comma:
996 self.write(", ")
997 else:
998 comma = True
999 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001000
1001 with self.block():
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001002 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001003
1004 def visit_FunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001005 self._function_helper(node, "def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001006
1007 def visit_AsyncFunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001008 self._function_helper(node, "async def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001009
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001010 def _function_helper(self, node, fill_suffix):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +03001011 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001012 for deco in node.decorator_list:
1013 self.fill("@")
1014 self.traverse(deco)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001015 def_str = fill_suffix + " " + node.name
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001016 self.fill(def_str)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001017 with self.delimit("(", ")"):
1018 self.traverse(node.args)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001019 if node.returns:
1020 self.write(" -> ")
1021 self.traverse(node.returns)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001022 with self.block(extra=self.get_type_comment(node)):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001023 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001024
1025 def visit_For(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001026 self._for_helper("for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001027
1028 def visit_AsyncFor(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001029 self._for_helper("async for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001030
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001031 def _for_helper(self, fill, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001032 self.fill(fill)
1033 self.traverse(node.target)
1034 self.write(" in ")
1035 self.traverse(node.iter)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001036 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001037 self.traverse(node.body)
1038 if node.orelse:
1039 self.fill("else")
1040 with self.block():
1041 self.traverse(node.orelse)
1042
1043 def visit_If(self, node):
1044 self.fill("if ")
1045 self.traverse(node.test)
1046 with self.block():
1047 self.traverse(node.body)
1048 # collapse nested ifs into equivalent elifs.
1049 while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], If):
1050 node = node.orelse[0]
1051 self.fill("elif ")
1052 self.traverse(node.test)
1053 with self.block():
1054 self.traverse(node.body)
1055 # final else
1056 if node.orelse:
1057 self.fill("else")
1058 with self.block():
1059 self.traverse(node.orelse)
1060
1061 def visit_While(self, node):
1062 self.fill("while ")
1063 self.traverse(node.test)
1064 with self.block():
1065 self.traverse(node.body)
1066 if node.orelse:
1067 self.fill("else")
1068 with self.block():
1069 self.traverse(node.orelse)
1070
1071 def visit_With(self, node):
1072 self.fill("with ")
1073 self.interleave(lambda: self.write(", "), self.traverse, node.items)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001074 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001075 self.traverse(node.body)
1076
1077 def visit_AsyncWith(self, node):
1078 self.fill("async with ")
1079 self.interleave(lambda: self.write(", "), self.traverse, node.items)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001080 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001081 self.traverse(node.body)
1082
Shantanua993e902020-11-20 13:16:42 -08001083 def _str_literal_helper(
1084 self, string, *, quote_types=_ALL_QUOTES, escape_special_whitespace=False
1085 ):
1086 """Helper for writing string literals, minimizing escapes.
1087 Returns the tuple (string literal to write, possible quote types).
1088 """
1089 def escape_char(c):
1090 # \n and \t are non-printable, but we only escape them if
1091 # escape_special_whitespace is True
1092 if not escape_special_whitespace and c in "\n\t":
1093 return c
1094 # Always escape backslashes and other non-printable characters
1095 if c == "\\" or not c.isprintable():
1096 return c.encode("unicode_escape").decode("ascii")
1097 return c
1098
1099 escaped_string = "".join(map(escape_char, string))
1100 possible_quotes = quote_types
1101 if "\n" in escaped_string:
1102 possible_quotes = [q for q in possible_quotes if q in _MULTI_QUOTES]
1103 possible_quotes = [q for q in possible_quotes if q not in escaped_string]
1104 if not possible_quotes:
1105 # If there aren't any possible_quotes, fallback to using repr
1106 # on the original string. Try to use a quote from quote_types,
1107 # e.g., so that we use triple quotes for docstrings.
1108 string = repr(string)
1109 quote = next((q for q in quote_types if string[0] in q), string[0])
1110 return string[1:-1], [quote]
1111 if escaped_string:
1112 # Sort so that we prefer '''"''' over """\""""
1113 possible_quotes.sort(key=lambda q: q[0] == escaped_string[-1])
1114 # If we're using triple quotes and we'd need to escape a final
1115 # quote, escape it
1116 if possible_quotes[0][0] == escaped_string[-1]:
1117 assert len(possible_quotes[0]) == 3
1118 escaped_string = escaped_string[:-1] + "\\" + escaped_string[-1]
1119 return escaped_string, possible_quotes
1120
1121 def _write_str_avoiding_backslashes(self, string, *, quote_types=_ALL_QUOTES):
1122 """Write string literal value with a best effort attempt to avoid backslashes."""
1123 string, quote_types = self._str_literal_helper(string, quote_types=quote_types)
1124 quote_type = quote_types[0]
1125 self.write(f"{quote_type}{string}{quote_type}")
1126
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001127 def visit_JoinedStr(self, node):
1128 self.write("f")
Shantanua993e902020-11-20 13:16:42 -08001129 if self._avoid_backslashes:
1130 self._fstring_JoinedStr(node, self.buffer_writer)
1131 self._write_str_avoiding_backslashes(self.buffer)
1132 return
1133
1134 # If we don't need to avoid backslashes globally (i.e., we only need
1135 # to avoid them inside FormattedValues), it's cosmetically preferred
1136 # to use escaped whitespace. That is, it's preferred to use backslashes
1137 # for cases like: f"{x}\n". To accomplish this, we keep track of what
1138 # in our buffer corresponds to FormattedValues and what corresponds to
1139 # Constant parts of the f-string, and allow escapes accordingly.
1140 buffer = []
1141 for value in node.values:
1142 meth = getattr(self, "_fstring_" + type(value).__name__)
1143 meth(value, self.buffer_writer)
1144 buffer.append((self.buffer, isinstance(value, Constant)))
1145 new_buffer = []
1146 quote_types = _ALL_QUOTES
1147 for value, is_constant in buffer:
1148 # Repeatedly narrow down the list of possible quote_types
1149 value, quote_types = self._str_literal_helper(
1150 value, quote_types=quote_types,
1151 escape_special_whitespace=is_constant
1152 )
1153 new_buffer.append(value)
1154 value = "".join(new_buffer)
1155 quote_type = quote_types[0]
1156 self.write(f"{quote_type}{value}{quote_type}")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001157
1158 def visit_FormattedValue(self, node):
1159 self.write("f")
1160 self._fstring_FormattedValue(node, self.buffer_writer)
Shantanua993e902020-11-20 13:16:42 -08001161 self._write_str_avoiding_backslashes(self.buffer)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001162
1163 def _fstring_JoinedStr(self, node, write):
1164 for value in node.values:
1165 meth = getattr(self, "_fstring_" + type(value).__name__)
1166 meth(value, write)
1167
1168 def _fstring_Constant(self, node, write):
1169 if not isinstance(node.value, str):
1170 raise ValueError("Constants inside JoinedStr should be a string.")
1171 value = node.value.replace("{", "{{").replace("}", "}}")
1172 write(value)
1173
1174 def _fstring_FormattedValue(self, node, write):
1175 write("{")
Shantanua993e902020-11-20 13:16:42 -08001176 unparser = type(self)(_avoid_backslashes=True)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001177 unparser.set_precedence(_Precedence.TEST.next(), node.value)
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +03001178 expr = unparser.visit(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001179 if expr.startswith("{"):
1180 write(" ") # Separate pair of opening brackets as "{ {"
Shantanua993e902020-11-20 13:16:42 -08001181 if "\\" in expr:
1182 raise ValueError("Unable to avoid backslash in f-string expression part")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001183 write(expr)
1184 if node.conversion != -1:
1185 conversion = chr(node.conversion)
1186 if conversion not in "sra":
1187 raise ValueError("Unknown f-string conversion.")
1188 write(f"!{conversion}")
1189 if node.format_spec:
1190 write(":")
1191 meth = getattr(self, "_fstring_" + type(node.format_spec).__name__)
1192 meth(node.format_spec, write)
1193 write("}")
1194
1195 def visit_Name(self, node):
1196 self.write(node.id)
1197
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001198 def _write_docstring(self, node):
1199 self.fill()
1200 if node.kind == "u":
1201 self.write("u")
Shantanua993e902020-11-20 13:16:42 -08001202 self._write_str_avoiding_backslashes(node.value, quote_types=_MULTI_QUOTES)
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001203
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001204 def _write_constant(self, value):
1205 if isinstance(value, (float, complex)):
Kodi Arfer08ff4362021-03-18 13:36:06 -04001206 # Substitute overflowing decimal literal for AST infinities,
1207 # and inf - inf for NaNs.
1208 self.write(
1209 repr(value)
1210 .replace("inf", _INFSTR)
1211 .replace("nan", f"({_INFSTR}-{_INFSTR})")
1212 )
Shantanua993e902020-11-20 13:16:42 -08001213 elif self._avoid_backslashes and isinstance(value, str):
1214 self._write_str_avoiding_backslashes(value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001215 else:
1216 self.write(repr(value))
1217
1218 def visit_Constant(self, node):
1219 value = node.value
1220 if isinstance(value, tuple):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001221 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001222 self.items_view(self._write_constant, value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001223 elif value is ...:
1224 self.write("...")
1225 else:
1226 if node.kind == "u":
1227 self.write("u")
1228 self._write_constant(node.value)
1229
1230 def visit_List(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001231 with self.delimit("[", "]"):
1232 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001233
1234 def visit_ListComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001235 with self.delimit("[", "]"):
1236 self.traverse(node.elt)
1237 for gen in node.generators:
1238 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001239
1240 def visit_GeneratorExp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001241 with self.delimit("(", ")"):
1242 self.traverse(node.elt)
1243 for gen in node.generators:
1244 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001245
1246 def visit_SetComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001247 with self.delimit("{", "}"):
1248 self.traverse(node.elt)
1249 for gen in node.generators:
1250 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001251
1252 def visit_DictComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001253 with self.delimit("{", "}"):
1254 self.traverse(node.key)
1255 self.write(": ")
1256 self.traverse(node.value)
1257 for gen in node.generators:
1258 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001259
1260 def visit_comprehension(self, node):
1261 if node.is_async:
1262 self.write(" async for ")
1263 else:
1264 self.write(" for ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001265 self.set_precedence(_Precedence.TUPLE, node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001266 self.traverse(node.target)
1267 self.write(" in ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001268 self.set_precedence(_Precedence.TEST.next(), node.iter, *node.ifs)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001269 self.traverse(node.iter)
1270 for if_clause in node.ifs:
1271 self.write(" if ")
1272 self.traverse(if_clause)
1273
1274 def visit_IfExp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001275 with self.require_parens(_Precedence.TEST, node):
1276 self.set_precedence(_Precedence.TEST.next(), node.body, node.test)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001277 self.traverse(node.body)
1278 self.write(" if ")
1279 self.traverse(node.test)
1280 self.write(" else ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001281 self.set_precedence(_Precedence.TEST, node.orelse)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001282 self.traverse(node.orelse)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001283
1284 def visit_Set(self, node):
Kodi Arfer08ff4362021-03-18 13:36:06 -04001285 if node.elts:
1286 with self.delimit("{", "}"):
1287 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
1288 else:
1289 # `{}` would be interpreted as a dictionary literal, and
1290 # `set` might be shadowed. Thus:
1291 self.write('{*()}')
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001292
1293 def visit_Dict(self, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001294 def write_key_value_pair(k, v):
1295 self.traverse(k)
1296 self.write(": ")
1297 self.traverse(v)
1298
1299 def write_item(item):
1300 k, v = item
1301 if k is None:
1302 # for dictionary unpacking operator in dicts {**{'y': 2}}
1303 # see PEP 448 for details
1304 self.write("**")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001305 self.set_precedence(_Precedence.EXPR, v)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001306 self.traverse(v)
1307 else:
1308 write_key_value_pair(k, v)
1309
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001310 with self.delimit("{", "}"):
1311 self.interleave(
1312 lambda: self.write(", "), write_item, zip(node.keys, node.values)
1313 )
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001314
1315 def visit_Tuple(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001316 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001317 self.items_view(self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001318
1319 unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001320 unop_precedence = {
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001321 "not": _Precedence.NOT,
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001322 "~": _Precedence.FACTOR,
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001323 "+": _Precedence.FACTOR,
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001324 "-": _Precedence.FACTOR,
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001325 }
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001326
1327 def visit_UnaryOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001328 operator = self.unop[node.op.__class__.__name__]
1329 operator_precedence = self.unop_precedence[operator]
1330 with self.require_parens(operator_precedence, node):
1331 self.write(operator)
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001332 # factor prefixes (+, -, ~) shouldn't be seperated
1333 # from the value they belong, (e.g: +1 instead of + 1)
1334 if operator_precedence is not _Precedence.FACTOR:
1335 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001336 self.set_precedence(operator_precedence, node.operand)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001337 self.traverse(node.operand)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001338
1339 binop = {
1340 "Add": "+",
1341 "Sub": "-",
1342 "Mult": "*",
1343 "MatMult": "@",
1344 "Div": "/",
1345 "Mod": "%",
1346 "LShift": "<<",
1347 "RShift": ">>",
1348 "BitOr": "|",
1349 "BitXor": "^",
1350 "BitAnd": "&",
1351 "FloorDiv": "//",
1352 "Pow": "**",
1353 }
1354
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001355 binop_precedence = {
1356 "+": _Precedence.ARITH,
1357 "-": _Precedence.ARITH,
1358 "*": _Precedence.TERM,
1359 "@": _Precedence.TERM,
1360 "/": _Precedence.TERM,
1361 "%": _Precedence.TERM,
1362 "<<": _Precedence.SHIFT,
1363 ">>": _Precedence.SHIFT,
1364 "|": _Precedence.BOR,
1365 "^": _Precedence.BXOR,
1366 "&": _Precedence.BAND,
1367 "//": _Precedence.TERM,
1368 "**": _Precedence.POWER,
1369 }
1370
1371 binop_rassoc = frozenset(("**",))
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001372 def visit_BinOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001373 operator = self.binop[node.op.__class__.__name__]
1374 operator_precedence = self.binop_precedence[operator]
1375 with self.require_parens(operator_precedence, node):
1376 if operator in self.binop_rassoc:
1377 left_precedence = operator_precedence.next()
1378 right_precedence = operator_precedence
1379 else:
1380 left_precedence = operator_precedence
1381 right_precedence = operator_precedence.next()
1382
1383 self.set_precedence(left_precedence, node.left)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001384 self.traverse(node.left)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001385 self.write(f" {operator} ")
1386 self.set_precedence(right_precedence, node.right)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001387 self.traverse(node.right)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001388
1389 cmpops = {
1390 "Eq": "==",
1391 "NotEq": "!=",
1392 "Lt": "<",
1393 "LtE": "<=",
1394 "Gt": ">",
1395 "GtE": ">=",
1396 "Is": "is",
1397 "IsNot": "is not",
1398 "In": "in",
1399 "NotIn": "not in",
1400 }
1401
1402 def visit_Compare(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001403 with self.require_parens(_Precedence.CMP, node):
1404 self.set_precedence(_Precedence.CMP.next(), node.left, *node.comparators)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001405 self.traverse(node.left)
1406 for o, e in zip(node.ops, node.comparators):
1407 self.write(" " + self.cmpops[o.__class__.__name__] + " ")
1408 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001409
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001410 boolops = {"And": "and", "Or": "or"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001411 boolop_precedence = {"and": _Precedence.AND, "or": _Precedence.OR}
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001412
1413 def visit_BoolOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001414 operator = self.boolops[node.op.__class__.__name__]
1415 operator_precedence = self.boolop_precedence[operator]
1416
1417 def increasing_level_traverse(node):
1418 nonlocal operator_precedence
1419 operator_precedence = operator_precedence.next()
1420 self.set_precedence(operator_precedence, node)
1421 self.traverse(node)
1422
1423 with self.require_parens(operator_precedence, node):
1424 s = f" {operator} "
1425 self.interleave(lambda: self.write(s), increasing_level_traverse, node.values)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001426
1427 def visit_Attribute(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001428 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001429 self.traverse(node.value)
1430 # Special case: 3.__abs__() is a syntax error, so if node.value
1431 # is an integer literal then we need to either parenthesize
1432 # it or add an extra space to get 3 .__abs__().
1433 if isinstance(node.value, Constant) and isinstance(node.value.value, int):
1434 self.write(" ")
1435 self.write(".")
1436 self.write(node.attr)
1437
1438 def visit_Call(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001439 self.set_precedence(_Precedence.ATOM, node.func)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001440 self.traverse(node.func)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001441 with self.delimit("(", ")"):
1442 comma = False
1443 for e in node.args:
1444 if comma:
1445 self.write(", ")
1446 else:
1447 comma = True
1448 self.traverse(e)
1449 for e in node.keywords:
1450 if comma:
1451 self.write(", ")
1452 else:
1453 comma = True
1454 self.traverse(e)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001455
1456 def visit_Subscript(self, node):
Batuhan Taskayac102a142020-05-18 23:48:49 +03001457 def is_simple_tuple(slice_value):
1458 # when unparsing a non-empty tuple, the parantheses can be safely
1459 # omitted if there aren't any elements that explicitly requires
1460 # parantheses (such as starred expressions).
1461 return (
1462 isinstance(slice_value, Tuple)
1463 and slice_value.elts
1464 and not any(isinstance(elt, Starred) for elt in slice_value.elts)
1465 )
1466
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001467 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001468 self.traverse(node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001469 with self.delimit("[", "]"):
Batuhan Taskayac102a142020-05-18 23:48:49 +03001470 if is_simple_tuple(node.slice):
Serhiy Storchaka13d52c22020-03-10 18:52:34 +02001471 self.items_view(self.traverse, node.slice.elts)
Serhiy Storchakac4928fc2020-03-07 17:25:32 +02001472 else:
1473 self.traverse(node.slice)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001474
1475 def visit_Starred(self, node):
1476 self.write("*")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001477 self.set_precedence(_Precedence.EXPR, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001478 self.traverse(node.value)
1479
1480 def visit_Ellipsis(self, node):
1481 self.write("...")
1482
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001483 def visit_Slice(self, node):
1484 if node.lower:
1485 self.traverse(node.lower)
1486 self.write(":")
1487 if node.upper:
1488 self.traverse(node.upper)
1489 if node.step:
1490 self.write(":")
1491 self.traverse(node.step)
1492
Brandt Bucher145bf262021-02-26 14:51:55 -08001493 def visit_Match(self, node):
1494 self.fill("match ")
1495 self.traverse(node.subject)
1496 with self.block():
1497 for case in node.cases:
1498 self.traverse(case)
1499
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001500 def visit_arg(self, node):
1501 self.write(node.arg)
1502 if node.annotation:
1503 self.write(": ")
1504 self.traverse(node.annotation)
1505
1506 def visit_arguments(self, node):
1507 first = True
1508 # normal arguments
1509 all_args = node.posonlyargs + node.args
1510 defaults = [None] * (len(all_args) - len(node.defaults)) + node.defaults
1511 for index, elements in enumerate(zip(all_args, defaults), 1):
1512 a, d = elements
1513 if first:
1514 first = False
1515 else:
1516 self.write(", ")
1517 self.traverse(a)
1518 if d:
1519 self.write("=")
1520 self.traverse(d)
1521 if index == len(node.posonlyargs):
1522 self.write(", /")
1523
1524 # varargs, or bare '*' if no varargs but keyword-only arguments present
1525 if node.vararg or node.kwonlyargs:
1526 if first:
1527 first = False
1528 else:
1529 self.write(", ")
1530 self.write("*")
1531 if node.vararg:
1532 self.write(node.vararg.arg)
1533 if node.vararg.annotation:
1534 self.write(": ")
1535 self.traverse(node.vararg.annotation)
1536
1537 # keyword-only arguments
1538 if node.kwonlyargs:
1539 for a, d in zip(node.kwonlyargs, node.kw_defaults):
Batuhan Taşkayaa322f502019-12-16 15:26:58 +03001540 self.write(", ")
1541 self.traverse(a)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001542 if d:
1543 self.write("=")
1544 self.traverse(d)
1545
1546 # kwargs
1547 if node.kwarg:
1548 if first:
1549 first = False
1550 else:
1551 self.write(", ")
1552 self.write("**" + node.kwarg.arg)
1553 if node.kwarg.annotation:
1554 self.write(": ")
1555 self.traverse(node.kwarg.annotation)
1556
1557 def visit_keyword(self, node):
1558 if node.arg is None:
1559 self.write("**")
1560 else:
1561 self.write(node.arg)
1562 self.write("=")
1563 self.traverse(node.value)
1564
1565 def visit_Lambda(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001566 with self.require_parens(_Precedence.TEST, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001567 self.write("lambda ")
1568 self.traverse(node.args)
1569 self.write(": ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001570 self.set_precedence(_Precedence.TEST, node.body)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001571 self.traverse(node.body)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001572
1573 def visit_alias(self, node):
1574 self.write(node.name)
1575 if node.asname:
1576 self.write(" as " + node.asname)
1577
1578 def visit_withitem(self, node):
1579 self.traverse(node.context_expr)
1580 if node.optional_vars:
1581 self.write(" as ")
1582 self.traverse(node.optional_vars)
1583
Brandt Bucher145bf262021-02-26 14:51:55 -08001584 def visit_match_case(self, node):
1585 self.fill("case ")
1586 self.traverse(node.pattern)
1587 if node.guard:
1588 self.write(" if ")
1589 self.traverse(node.guard)
1590 with self.block():
1591 self.traverse(node.body)
1592
Nick Coghlan1e7b8582021-04-29 15:58:44 +10001593 def visit_MatchValue(self, node):
1594 self.traverse(node.value)
1595
1596 def visit_MatchSingleton(self, node):
1597 self._write_constant(node.value)
1598
1599 def visit_MatchSequence(self, node):
1600 with self.delimit("[", "]"):
1601 self.interleave(
1602 lambda: self.write(", "), self.traverse, node.patterns
1603 )
1604
1605 def visit_MatchStar(self, node):
1606 name = node.name
1607 if name is None:
1608 name = "_"
1609 self.write(f"*{name}")
1610
1611 def visit_MatchMapping(self, node):
1612 def write_key_pattern_pair(pair):
1613 k, p = pair
1614 self.traverse(k)
1615 self.write(": ")
1616 self.traverse(p)
1617
1618 with self.delimit("{", "}"):
1619 keys = node.keys
1620 self.interleave(
1621 lambda: self.write(", "),
1622 write_key_pattern_pair,
1623 zip(keys, node.patterns, strict=True),
1624 )
1625 rest = node.rest
1626 if rest is not None:
1627 if keys:
1628 self.write(", ")
1629 self.write(f"**{rest}")
1630
1631 def visit_MatchClass(self, node):
1632 self.set_precedence(_Precedence.ATOM, node.cls)
1633 self.traverse(node.cls)
1634 with self.delimit("(", ")"):
1635 patterns = node.patterns
1636 self.interleave(
1637 lambda: self.write(", "), self.traverse, patterns
1638 )
1639 attrs = node.kwd_attrs
1640 if attrs:
1641 def write_attr_pattern(pair):
1642 attr, pattern = pair
1643 self.write(f"{attr}=")
1644 self.traverse(pattern)
1645
1646 if patterns:
1647 self.write(", ")
1648 self.interleave(
1649 lambda: self.write(", "),
1650 write_attr_pattern,
1651 zip(attrs, node.kwd_patterns, strict=True),
1652 )
1653
Brandt Bucher145bf262021-02-26 14:51:55 -08001654 def visit_MatchAs(self, node):
Nick Coghlan1e7b8582021-04-29 15:58:44 +10001655 name = node.name
1656 pattern = node.pattern
1657 if name is None:
1658 self.write("_")
1659 elif pattern is None:
1660 self.write(node.name)
1661 else:
1662 with self.require_parens(_Precedence.TEST, node):
1663 self.set_precedence(_Precedence.BOR, node.pattern)
1664 self.traverse(node.pattern)
1665 self.write(f" as {node.name}")
Brandt Bucher145bf262021-02-26 14:51:55 -08001666
1667 def visit_MatchOr(self, node):
1668 with self.require_parens(_Precedence.BOR, node):
1669 self.set_precedence(_Precedence.BOR.next(), *node.patterns)
1670 self.interleave(lambda: self.write(" | "), self.traverse, node.patterns)
1671
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001672def unparse(ast_obj):
1673 unparser = _Unparser()
1674 return unparser.visit(ast_obj)
1675
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001676
1677def main():
1678 import argparse
1679
1680 parser = argparse.ArgumentParser(prog='python -m ast')
1681 parser.add_argument('infile', type=argparse.FileType(mode='rb'), nargs='?',
1682 default='-',
1683 help='the file to parse; defaults to stdin')
1684 parser.add_argument('-m', '--mode', default='exec',
1685 choices=('exec', 'single', 'eval', 'func_type'),
1686 help='specify what kind of code must be parsed')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001687 parser.add_argument('--no-type-comments', default=True, action='store_false',
1688 help="don't add information about type comments")
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001689 parser.add_argument('-a', '--include-attributes', action='store_true',
1690 help='include attributes such as line numbers and '
1691 'column offsets')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001692 parser.add_argument('-i', '--indent', type=int, default=3,
1693 help='indentation of nodes (number of spaces)')
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001694 args = parser.parse_args()
1695
1696 with args.infile as infile:
1697 source = infile.read()
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001698 tree = parse(source, args.infile.name, args.mode, type_comments=args.no_type_comments)
1699 print(dump(tree, include_attributes=args.include_attributes, indent=args.indent))
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001700
1701if __name__ == '__main__':
1702 main()