blob: f4d2f6e42c022be5fad3167206c26758f99ee844 [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 Furman9bf7c2d2021-07-03 21:08:42 -070030from enum import IntEnum, auto
Georg Brandl0c77a822008-06-10 16:37:50 +000031
32
Guido van Rossum495da292019-03-07 12:38:08 -080033def parse(source, filename='<unknown>', mode='exec', *,
Guido van Rossum10b55c12019-06-11 17:23:12 -070034 type_comments=False, feature_version=None):
Georg Brandl0c77a822008-06-10 16:37:50 +000035 """
Terry Reedyfeac6242011-01-24 21:36:03 +000036 Parse the source into an AST node.
37 Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
Guido van Rossumdcfcd142019-01-31 03:40:27 -080038 Pass type_comments=True to get back type comments where the syntax allows.
Georg Brandl0c77a822008-06-10 16:37:50 +000039 """
Guido van Rossumdcfcd142019-01-31 03:40:27 -080040 flags = PyCF_ONLY_AST
41 if type_comments:
42 flags |= PyCF_TYPE_COMMENTS
Guido van Rossum10b55c12019-06-11 17:23:12 -070043 if isinstance(feature_version, tuple):
44 major, minor = feature_version # Should be a 2-tuple.
45 assert major == 3
46 feature_version = minor
47 elif feature_version is None:
48 feature_version = -1
49 # Else it should be an int giving the minor version for 3.x.
Guido van Rossum495da292019-03-07 12:38:08 -080050 return compile(source, filename, mode, flags,
Victor Stinnerefdf6ca2019-06-12 02:52:16 +020051 _feature_version=feature_version)
Georg Brandl0c77a822008-06-10 16:37:50 +000052
53
54def literal_eval(node_or_string):
55 """
56 Safely evaluate an expression node or a string containing a Python
57 expression. The string or node provided may only consist of the following
Éric Araujo2a83cc62011-04-17 19:10:27 +020058 Python literal structures: strings, bytes, numbers, tuples, lists, dicts,
59 sets, booleans, and None.
Georg Brandl0c77a822008-06-10 16:37:50 +000060 """
Georg Brandl0c77a822008-06-10 16:37:50 +000061 if isinstance(node_or_string, str):
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 Furman9bf7c2d2021-07-03 21:08:42 -0700639class _Precedence(IntEnum):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300640 """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
Shantanua993e902020-11-20 13:16:42 -0800668
669_SINGLE_QUOTES = ("'", '"')
670_MULTI_QUOTES = ('"""', "'''")
671_ALL_QUOTES = (*_SINGLE_QUOTES, *_MULTI_QUOTES)
672
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000673class _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
Shantanua993e902020-11-20 13:16:42 -0800678 def __init__(self, *, _avoid_backslashes=False):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000679 self._source = []
680 self._buffer = []
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300681 self._precedences = {}
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300682 self._type_ignores = {}
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000683 self._indent = 0
Shantanua993e902020-11-20 13:16:42 -0800684 self._avoid_backslashes = _avoid_backslashes
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000685
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
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +0300698 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
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300708 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
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000713 def fill(self, text=""):
714 """Indent a piece of text and append it, according to the current
715 indentation level"""
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300716 self.maybe_newline()
717 self.write(" " * self._indent + text)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000718
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
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000732 @contextmanager
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300733 def block(self, *, extra = None):
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000734 """A context manager for preparing the source for blocks. It adds
735 the character':', increases the indentation on enter and decreases
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300736 the indentation on exit. If *extra* is given, it will be directly
737 appended after the colon character.
738 """
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000739 self.write(":")
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300740 if extra:
741 self.write(extra)
Pablo Galindod69cbeb2019-12-23 16:42:48 +0000742 self._indent += 1
743 yield
744 self._indent -= 1
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000745
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300746 @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
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300761 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
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300772 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
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300788 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
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000793 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
Nick Coghlan1e7b8582021-04-29 15:58:44 +1000800 # 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000803 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
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300810 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
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000817 def visit_Module(self, node):
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300818 self._type_ignores = {
819 ignore.lineno: f"ignore{ignore.tag}"
820 for ignore in node.type_ignores
821 }
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +0300822 self._write_docstring_and_traverse_body(node)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300823 self._type_ignores.clear()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000824
Batuhan Taşkaya5b66ec12020-03-15 22:56:57 +0300825 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
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000834 def visit_Expr(self, node):
835 self.fill()
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300836 self.set_precedence(_Precedence.YIELD, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000837 self.traverse(node.value)
838
839 def visit_NamedExpr(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300840 with self.require_parens(_Precedence.TUPLE, node):
841 self.set_precedence(_Precedence.ATOM, node.target, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300842 self.traverse(node.target)
843 self.write(" := ")
844 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000845
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)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +0300864 if type_comment := self.get_type_comment(node):
865 self.write(type_comment)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000866
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()
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300875 with self.delimit_if("(", ")", not node.simple and isinstance(node.target, Name)):
876 self.traverse(node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000877 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):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300918 with self.require_parens(_Precedence.AWAIT, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300919 self.write("await")
920 if node.value:
921 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300922 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300923 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000924
925 def visit_Yield(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300926 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300927 self.write("yield")
928 if node.value:
929 self.write(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300930 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300931 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000932
933 def visit_YieldFrom(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300934 with self.require_parens(_Precedence.YIELD, node):
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300935 self.write("yield from ")
936 if not node.value:
937 raise ValueError("Node can't be used without a value attribute.")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +0300938 self.set_precedence(_Precedence.ATOM, node.value)
Batuhan Taşkaya7b35bef2020-01-02 21:20:04 +0300939 self.traverse(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000940
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):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +0300980 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000981 for deco in node.decorator_list:
982 self.fill("@")
983 self.traverse(deco)
984 self.fill("class " + node.name)
Batuhan Taskaya25160cd2020-05-17 00:53:25 +0300985 with self.delimit_if("(", ")", condition = node.bases or node.keywords):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +0300986 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +0000999
1000 with self.block():
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001001 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001002
1003 def visit_FunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001004 self._function_helper(node, "def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001005
1006 def visit_AsyncFunctionDef(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001007 self._function_helper(node, "async def")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001008
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001009 def _function_helper(self, node, fill_suffix):
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +03001010 self.maybe_newline()
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001011 for deco in node.decorator_list:
1012 self.fill("@")
1013 self.traverse(deco)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001014 def_str = fill_suffix + " " + node.name
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001015 self.fill(def_str)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001016 with self.delimit("(", ")"):
1017 self.traverse(node.args)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001018 if node.returns:
1019 self.write(" -> ")
1020 self.traverse(node.returns)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001021 with self.block(extra=self.get_type_comment(node)):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001022 self._write_docstring_and_traverse_body(node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001023
1024 def visit_For(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001025 self._for_helper("for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001026
1027 def visit_AsyncFor(self, node):
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001028 self._for_helper("async for ", node)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001029
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001030 def _for_helper(self, fill, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001031 self.fill(fill)
1032 self.traverse(node.target)
1033 self.write(" in ")
1034 self.traverse(node.iter)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001035 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001036 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)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001073 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001074 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)
Batuhan Taskayadff92bb2020-05-17 02:04:12 +03001079 with self.block(extra=self.get_type_comment(node)):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001080 self.traverse(node.body)
1081
Shantanua993e902020-11-20 13:16:42 -08001082 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
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001126 def visit_JoinedStr(self, node):
1127 self.write("f")
Shantanua993e902020-11-20 13:16:42 -08001128 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}")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001156
1157 def visit_FormattedValue(self, node):
1158 self.write("f")
1159 self._fstring_FormattedValue(node, self.buffer_writer)
Shantanua993e902020-11-20 13:16:42 -08001160 self._write_str_avoiding_backslashes(self.buffer)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001161
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("{")
Shantanua993e902020-11-20 13:16:42 -08001175 unparser = type(self)(_avoid_backslashes=True)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001176 unparser.set_precedence(_Precedence.TEST.next(), node.value)
Batuhan Taskaya493bf1c2020-05-03 20:11:51 +03001177 expr = unparser.visit(node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001178 if expr.startswith("{"):
1179 write(" ") # Separate pair of opening brackets as "{ {"
Shantanua993e902020-11-20 13:16:42 -08001180 if "\\" in expr:
1181 raise ValueError("Unable to avoid backslash in f-string expression part")
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001182 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
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001197 def _write_docstring(self, node):
1198 self.fill()
1199 if node.kind == "u":
1200 self.write("u")
Shantanua993e902020-11-20 13:16:42 -08001201 self._write_str_avoiding_backslashes(node.value, quote_types=_MULTI_QUOTES)
Batuhan Taşkaya89aa4692020-03-02 21:59:01 +03001202
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001203 def _write_constant(self, value):
1204 if isinstance(value, (float, complex)):
Kodi Arfer08ff4362021-03-18 13:36:06 -04001205 # 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 )
Shantanua993e902020-11-20 13:16:42 -08001212 elif self._avoid_backslashes and isinstance(value, str):
1213 self._write_str_avoiding_backslashes(value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001214 else:
1215 self.write(repr(value))
1216
1217 def visit_Constant(self, node):
1218 value = node.value
1219 if isinstance(value, tuple):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001220 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001221 self.items_view(self._write_constant, value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001222 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):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001230 with self.delimit("[", "]"):
1231 self.interleave(lambda: self.write(", "), self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001232
1233 def visit_ListComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001234 with self.delimit("[", "]"):
1235 self.traverse(node.elt)
1236 for gen in node.generators:
1237 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001238
1239 def visit_GeneratorExp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001240 with self.delimit("(", ")"):
1241 self.traverse(node.elt)
1242 for gen in node.generators:
1243 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001244
1245 def visit_SetComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001246 with self.delimit("{", "}"):
1247 self.traverse(node.elt)
1248 for gen in node.generators:
1249 self.traverse(gen)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001250
1251 def visit_DictComp(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001252 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001258
1259 def visit_comprehension(self, node):
1260 if node.is_async:
1261 self.write(" async for ")
1262 else:
1263 self.write(" for ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001264 self.set_precedence(_Precedence.TUPLE, node.target)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001265 self.traverse(node.target)
1266 self.write(" in ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001267 self.set_precedence(_Precedence.TEST.next(), node.iter, *node.ifs)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001268 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):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001274 with self.require_parens(_Precedence.TEST, node):
1275 self.set_precedence(_Precedence.TEST.next(), node.body, node.test)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001276 self.traverse(node.body)
1277 self.write(" if ")
1278 self.traverse(node.test)
1279 self.write(" else ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001280 self.set_precedence(_Precedence.TEST, node.orelse)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001281 self.traverse(node.orelse)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001282
1283 def visit_Set(self, node):
Kodi Arfer08ff4362021-03-18 13:36:06 -04001284 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('{*()}')
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001291
1292 def visit_Dict(self, node):
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001293 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("**")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001304 self.set_precedence(_Precedence.EXPR, v)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001305 self.traverse(v)
1306 else:
1307 write_key_value_pair(k, v)
1308
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001309 with self.delimit("{", "}"):
1310 self.interleave(
1311 lambda: self.write(", "), write_item, zip(node.keys, node.values)
1312 )
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001313
1314 def visit_Tuple(self, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001315 with self.delimit("(", ")"):
Batuhan Taşkayae7cab7f2020-03-09 23:27:03 +03001316 self.items_view(self.traverse, node.elts)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001317
1318 unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001319 unop_precedence = {
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001320 "not": _Precedence.NOT,
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001321 "~": _Precedence.FACTOR,
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001322 "+": _Precedence.FACTOR,
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001323 "-": _Precedence.FACTOR,
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001324 }
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001325
1326 def visit_UnaryOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001327 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)
Batuhan Taskayace4a7532020-05-17 00:46:11 +03001331 # 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(" ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001335 self.set_precedence(operator_precedence, node.operand)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001336 self.traverse(node.operand)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001337
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
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001354 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(("**",))
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001371 def visit_BinOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001372 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)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001383 self.traverse(node.left)
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001384 self.write(f" {operator} ")
1385 self.set_precedence(right_precedence, node.right)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001386 self.traverse(node.right)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001387
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):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001402 with self.require_parens(_Precedence.CMP, node):
1403 self.set_precedence(_Precedence.CMP.next(), node.left, *node.comparators)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001404 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001408
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001409 boolops = {"And": "and", "Or": "or"}
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001410 boolop_precedence = {"and": _Precedence.AND, "or": _Precedence.OR}
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001411
1412 def visit_BoolOp(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001413 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001425
1426 def visit_Attribute(self, node):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001427 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001428 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):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001438 self.set_precedence(_Precedence.ATOM, node.func)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001439 self.traverse(node.func)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001440 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)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001454
1455 def visit_Subscript(self, node):
Batuhan Taskayac102a142020-05-18 23:48:49 +03001456 def is_simple_tuple(slice_value):
Miss Islington (bot)62a983d2021-05-04 06:39:16 -07001457 # when unparsing a non-empty tuple, the parentheses can be safely
Batuhan Taskayac102a142020-05-18 23:48:49 +03001458 # omitted if there aren't any elements that explicitly requires
Miss Islington (bot)62a983d2021-05-04 06:39:16 -07001459 # parentheses (such as starred expressions).
Batuhan Taskayac102a142020-05-18 23:48:49 +03001460 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
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001466 self.set_precedence(_Precedence.ATOM, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001467 self.traverse(node.value)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001468 with self.delimit("[", "]"):
Batuhan Taskayac102a142020-05-18 23:48:49 +03001469 if is_simple_tuple(node.slice):
Serhiy Storchaka13d52c22020-03-10 18:52:34 +02001470 self.items_view(self.traverse, node.slice.elts)
Serhiy Storchakac4928fc2020-03-07 17:25:32 +02001471 else:
1472 self.traverse(node.slice)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001473
1474 def visit_Starred(self, node):
1475 self.write("*")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001476 self.set_precedence(_Precedence.EXPR, node.value)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001477 self.traverse(node.value)
1478
1479 def visit_Ellipsis(self, node):
1480 self.write("...")
1481
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001482 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
Brandt Bucher145bf262021-02-26 14:51:55 -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
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001499 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):
Batuhan Taşkayaa322f502019-12-16 15:26:58 +03001539 self.write(", ")
1540 self.traverse(a)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001541 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):
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001565 with self.require_parens(_Precedence.TEST, node):
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001566 self.write("lambda ")
1567 self.traverse(node.args)
1568 self.write(": ")
Batuhan Taşkaya397b96f2020-03-01 23:12:17 +03001569 self.set_precedence(_Precedence.TEST, node.body)
Batuhan Taşkaya4b3b1222019-12-23 19:11:00 +03001570 self.traverse(node.body)
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001571
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
Brandt Bucher145bf262021-02-26 14:51:55 -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
Nick Coghlan1e7b8582021-04-29 15:58:44 +10001592 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
Brandt Bucher145bf262021-02-26 14:51:55 -08001653 def visit_MatchAs(self, node):
Nick Coghlan1e7b8582021-04-29 15:58:44 +10001654 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}")
Brandt Bucher145bf262021-02-26 14:51:55 -08001665
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
Pablo Galindo27fc3b62019-11-24 23:02:40 +00001671def unparse(ast_obj):
1672 unparser = _Unparser()
1673 return unparser.visit(ast_obj)
1674
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001675
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')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001686 parser.add_argument('--no-type-comments', default=True, action='store_false',
1687 help="don't add information about type comments")
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001688 parser.add_argument('-a', '--include-attributes', action='store_true',
1689 help='include attributes such as line numbers and '
1690 'column offsets')
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001691 parser.add_argument('-i', '--indent', type=int, default=3,
1692 help='indentation of nodes (number of spaces)')
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001693 args = parser.parse_args()
1694
1695 with args.infile as infile:
1696 source = infile.read()
Batuhan Taşkaya814d6872019-12-16 21:23:27 +03001697 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))
Serhiy Storchaka832e8642019-09-09 23:36:13 +03001699
1700if __name__ == '__main__':
1701 main()