blob: 6955a7ada50cf6748f6531fb63d67e4a229dea83 [file] [log] [blame]
Patrick Williamsc124f4f2015-09-15 14:41:29 -05001# -*- coding: utf-8 -*-
2"""
3 codegen
4 ~~~~~~~
5
6 Extension to ast that allow ast -> python code generation.
7
8 :copyright: Copyright 2008 by Armin Ronacher.
9 :license: BSD.
10"""
11from ast import *
12
13BOOLOP_SYMBOLS = {
14 And: 'and',
15 Or: 'or'
16}
17
18BINOP_SYMBOLS = {
19 Add: '+',
20 Sub: '-',
21 Mult: '*',
22 Div: '/',
23 FloorDiv: '//',
24 Mod: '%',
25 LShift: '<<',
26 RShift: '>>',
27 BitOr: '|',
28 BitAnd: '&',
29 BitXor: '^'
30}
31
32CMPOP_SYMBOLS = {
33 Eq: '==',
34 Gt: '>',
35 GtE: '>=',
36 In: 'in',
37 Is: 'is',
38 IsNot: 'is not',
39 Lt: '<',
40 LtE: '<=',
41 NotEq: '!=',
42 NotIn: 'not in'
43}
44
45UNARYOP_SYMBOLS = {
46 Invert: '~',
47 Not: 'not',
48 UAdd: '+',
49 USub: '-'
50}
51
52ALL_SYMBOLS = {}
53ALL_SYMBOLS.update(BOOLOP_SYMBOLS)
54ALL_SYMBOLS.update(BINOP_SYMBOLS)
55ALL_SYMBOLS.update(CMPOP_SYMBOLS)
56ALL_SYMBOLS.update(UNARYOP_SYMBOLS)
57
58def to_source(node, indent_with=' ' * 4, add_line_information=False):
59 """This function can convert a node tree back into python sourcecode.
60 This is useful for debugging purposes, especially if you're dealing with
61 custom asts not generated by python itself.
62
63 It could be that the sourcecode is evaluable when the AST itself is not
64 compilable / evaluable. The reason for this is that the AST contains some
65 more data than regular sourcecode does, which is dropped during
66 conversion.
67
68 Each level of indentation is replaced with `indent_with`. Per default this
69 parameter is equal to four spaces as suggested by PEP 8, but it might be
70 adjusted to match the application's styleguide.
71
72 If `add_line_information` is set to `True` comments for the line numbers
73 of the nodes are added to the output. This can be used to spot wrong line
74 number information of statement nodes.
75 """
76 generator = SourceGenerator(indent_with, add_line_information)
77 generator.visit(node)
78 return ''.join(generator.result)
79
80
81class SourceGenerator(NodeVisitor):
82 """This visitor is able to transform a well formed syntax tree into python
83 sourcecode. For more details have a look at the docstring of the
84 `node_to_source` function.
85 """
86
87 def __init__(self, indent_with, add_line_information=False):
88 self.result = []
89 self.indent_with = indent_with
90 self.add_line_information = add_line_information
91 self.indentation = 0
92 self.new_lines = 0
93
94 def write(self, x):
95 if self.new_lines:
96 if self.result:
97 self.result.append('\n' * self.new_lines)
98 self.result.append(self.indent_with * self.indentation)
99 self.new_lines = 0
100 self.result.append(x)
101
102 def newline(self, node=None, extra=0):
103 self.new_lines = max(self.new_lines, 1 + extra)
104 if node is not None and self.add_line_information:
105 self.write('# line: %s' % node.lineno)
106 self.new_lines = 1
107
108 def body(self, statements):
109 self.new_line = True
110 self.indentation += 1
111 for stmt in statements:
112 self.visit(stmt)
113 self.indentation -= 1
114
115 def body_or_else(self, node):
116 self.body(node.body)
117 if node.orelse:
118 self.newline()
119 self.write('else:')
120 self.body(node.orelse)
121
122 def signature(self, node):
123 want_comma = []
124 def write_comma():
125 if want_comma:
126 self.write(', ')
127 else:
128 want_comma.append(True)
129
130 padding = [None] * (len(node.args) - len(node.defaults))
131 for arg, default in zip(node.args, padding + node.defaults):
132 write_comma()
133 self.visit(arg)
134 if default is not None:
135 self.write('=')
136 self.visit(default)
137 if node.vararg is not None:
138 write_comma()
139 self.write('*' + node.vararg)
140 if node.kwarg is not None:
141 write_comma()
142 self.write('**' + node.kwarg)
143
144 def decorators(self, node):
145 for decorator in node.decorator_list:
146 self.newline(decorator)
147 self.write('@')
148 self.visit(decorator)
149
150 # Statements
151
152 def visit_Assign(self, node):
153 self.newline(node)
154 for idx, target in enumerate(node.targets):
155 if idx:
156 self.write(', ')
157 self.visit(target)
158 self.write(' = ')
159 self.visit(node.value)
160
161 def visit_AugAssign(self, node):
162 self.newline(node)
163 self.visit(node.target)
164 self.write(BINOP_SYMBOLS[type(node.op)] + '=')
165 self.visit(node.value)
166
167 def visit_ImportFrom(self, node):
168 self.newline(node)
169 self.write('from %s%s import ' % ('.' * node.level, node.module))
170 for idx, item in enumerate(node.names):
171 if idx:
172 self.write(', ')
173 self.write(item)
174
175 def visit_Import(self, node):
176 self.newline(node)
177 for item in node.names:
178 self.write('import ')
179 self.visit(item)
180
181 def visit_Expr(self, node):
182 self.newline(node)
183 self.generic_visit(node)
184
185 def visit_FunctionDef(self, node):
186 self.newline(extra=1)
187 self.decorators(node)
188 self.newline(node)
189 self.write('def %s(' % node.name)
190 self.signature(node.args)
191 self.write('):')
192 self.body(node.body)
193
194 def visit_ClassDef(self, node):
195 have_args = []
196 def paren_or_comma():
197 if have_args:
198 self.write(', ')
199 else:
200 have_args.append(True)
201 self.write('(')
202
203 self.newline(extra=2)
204 self.decorators(node)
205 self.newline(node)
206 self.write('class %s' % node.name)
207 for base in node.bases:
208 paren_or_comma()
209 self.visit(base)
210 # XXX: the if here is used to keep this module compatible
211 # with python 2.6.
212 if hasattr(node, 'keywords'):
213 for keyword in node.keywords:
214 paren_or_comma()
215 self.write(keyword.arg + '=')
216 self.visit(keyword.value)
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600217 if hasattr(node, 'starargs') and node.starargs is not None:
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500218 paren_or_comma()
219 self.write('*')
220 self.visit(node.starargs)
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600221 if hasattr(node, 'kwargs') and node.kwargs is not None:
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500222 paren_or_comma()
223 self.write('**')
224 self.visit(node.kwargs)
225 self.write(have_args and '):' or ':')
226 self.body(node.body)
227
228 def visit_If(self, node):
229 self.newline(node)
230 self.write('if ')
231 self.visit(node.test)
232 self.write(':')
233 self.body(node.body)
234 while True:
235 else_ = node.orelse
236 if len(else_) == 1 and isinstance(else_[0], If):
237 node = else_[0]
238 self.newline()
239 self.write('elif ')
240 self.visit(node.test)
241 self.write(':')
242 self.body(node.body)
243 else:
244 self.newline()
245 self.write('else:')
246 self.body(else_)
247 break
248
249 def visit_For(self, node):
250 self.newline(node)
251 self.write('for ')
252 self.visit(node.target)
253 self.write(' in ')
254 self.visit(node.iter)
255 self.write(':')
256 self.body_or_else(node)
257
258 def visit_While(self, node):
259 self.newline(node)
260 self.write('while ')
261 self.visit(node.test)
262 self.write(':')
263 self.body_or_else(node)
264
265 def visit_With(self, node):
266 self.newline(node)
267 self.write('with ')
268 self.visit(node.context_expr)
269 if node.optional_vars is not None:
270 self.write(' as ')
271 self.visit(node.optional_vars)
272 self.write(':')
273 self.body(node.body)
274
275 def visit_Pass(self, node):
276 self.newline(node)
277 self.write('pass')
278
279 def visit_Print(self, node):
280 # XXX: python 2.6 only
281 self.newline(node)
282 self.write('print ')
283 want_comma = False
284 if node.dest is not None:
285 self.write(' >> ')
286 self.visit(node.dest)
287 want_comma = True
288 for value in node.values:
289 if want_comma:
290 self.write(', ')
291 self.visit(value)
292 want_comma = True
293 if not node.nl:
294 self.write(',')
295
296 def visit_Delete(self, node):
297 self.newline(node)
298 self.write('del ')
299 for idx, target in enumerate(node):
300 if idx:
301 self.write(', ')
302 self.visit(target)
303
304 def visit_TryExcept(self, node):
305 self.newline(node)
306 self.write('try:')
307 self.body(node.body)
308 for handler in node.handlers:
309 self.visit(handler)
310
311 def visit_TryFinally(self, node):
312 self.newline(node)
313 self.write('try:')
314 self.body(node.body)
315 self.newline(node)
316 self.write('finally:')
317 self.body(node.finalbody)
318
319 def visit_Global(self, node):
320 self.newline(node)
321 self.write('global ' + ', '.join(node.names))
322
323 def visit_Nonlocal(self, node):
324 self.newline(node)
325 self.write('nonlocal ' + ', '.join(node.names))
326
327 def visit_Return(self, node):
328 self.newline(node)
329 self.write('return ')
330 self.visit(node.value)
331
332 def visit_Break(self, node):
333 self.newline(node)
334 self.write('break')
335
336 def visit_Continue(self, node):
337 self.newline(node)
338 self.write('continue')
339
340 def visit_Raise(self, node):
341 # XXX: Python 2.6 / 3.0 compatibility
342 self.newline(node)
343 self.write('raise')
344 if hasattr(node, 'exc') and node.exc is not None:
345 self.write(' ')
346 self.visit(node.exc)
347 if node.cause is not None:
348 self.write(' from ')
349 self.visit(node.cause)
350 elif hasattr(node, 'type') and node.type is not None:
351 self.visit(node.type)
352 if node.inst is not None:
353 self.write(', ')
354 self.visit(node.inst)
355 if node.tback is not None:
356 self.write(', ')
357 self.visit(node.tback)
358
359 # Expressions
360
361 def visit_Attribute(self, node):
362 self.visit(node.value)
363 self.write('.' + node.attr)
364
365 def visit_Call(self, node):
366 want_comma = []
367 def write_comma():
368 if want_comma:
369 self.write(', ')
370 else:
371 want_comma.append(True)
372
373 self.visit(node.func)
374 self.write('(')
375 for arg in node.args:
376 write_comma()
377 self.visit(arg)
378 for keyword in node.keywords:
379 write_comma()
380 self.write(keyword.arg + '=')
381 self.visit(keyword.value)
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600382 if hasattr(node, 'starargs') and node.starargs is not None:
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500383 write_comma()
384 self.write('*')
385 self.visit(node.starargs)
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600386 if hasattr(node, 'kwargs') and node.kwargs is not None:
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500387 write_comma()
388 self.write('**')
389 self.visit(node.kwargs)
390 self.write(')')
391
392 def visit_Name(self, node):
393 self.write(node.id)
394
395 def visit_Str(self, node):
396 self.write(repr(node.s))
397
398 def visit_Bytes(self, node):
399 self.write(repr(node.s))
400
401 def visit_Num(self, node):
402 self.write(repr(node.n))
403
Andrew Geisslereff27472021-10-29 15:35:00 -0500404 def visit_Constant(self, node):
405 # Python 3.8 deprecated visit_Num(), visit_Str(), visit_Bytes(),
406 # visit_NameConstant() and visit_Ellipsis(). They can be removed once we
407 # require 3.8+.
408 self.write(repr(node.value))
409
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500410 def visit_Tuple(self, node):
411 self.write('(')
412 idx = -1
413 for idx, item in enumerate(node.elts):
414 if idx:
415 self.write(', ')
416 self.visit(item)
417 self.write(idx and ')' or ',)')
418
419 def sequence_visit(left, right):
420 def visit(self, node):
421 self.write(left)
422 for idx, item in enumerate(node.elts):
423 if idx:
424 self.write(', ')
425 self.visit(item)
426 self.write(right)
427 return visit
428
429 visit_List = sequence_visit('[', ']')
430 visit_Set = sequence_visit('{', '}')
431 del sequence_visit
432
433 def visit_Dict(self, node):
434 self.write('{')
435 for idx, (key, value) in enumerate(zip(node.keys, node.values)):
436 if idx:
437 self.write(', ')
438 self.visit(key)
439 self.write(': ')
440 self.visit(value)
441 self.write('}')
442
443 def visit_BinOp(self, node):
444 self.visit(node.left)
445 self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
446 self.visit(node.right)
447
448 def visit_BoolOp(self, node):
449 self.write('(')
450 for idx, value in enumerate(node.values):
451 if idx:
452 self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
453 self.visit(value)
454 self.write(')')
455
456 def visit_Compare(self, node):
457 self.write('(')
458 self.write(node.left)
459 for op, right in zip(node.ops, node.comparators):
460 self.write(' %s %%' % CMPOP_SYMBOLS[type(op)])
461 self.visit(right)
462 self.write(')')
463
464 def visit_UnaryOp(self, node):
465 self.write('(')
466 op = UNARYOP_SYMBOLS[type(node.op)]
467 self.write(op)
468 if op == 'not':
469 self.write(' ')
470 self.visit(node.operand)
471 self.write(')')
472
473 def visit_Subscript(self, node):
474 self.visit(node.value)
475 self.write('[')
476 self.visit(node.slice)
477 self.write(']')
478
479 def visit_Slice(self, node):
480 if node.lower is not None:
481 self.visit(node.lower)
482 self.write(':')
483 if node.upper is not None:
484 self.visit(node.upper)
485 if node.step is not None:
486 self.write(':')
487 if not (isinstance(node.step, Name) and node.step.id == 'None'):
488 self.visit(node.step)
489
490 def visit_ExtSlice(self, node):
491 for idx, item in node.dims:
492 if idx:
493 self.write(', ')
494 self.visit(item)
495
496 def visit_Yield(self, node):
497 self.write('yield ')
498 self.visit(node.value)
499
500 def visit_Lambda(self, node):
501 self.write('lambda ')
502 self.signature(node.args)
503 self.write(': ')
504 self.visit(node.body)
505
506 def visit_Ellipsis(self, node):
507 self.write('Ellipsis')
508
509 def generator_visit(left, right):
510 def visit(self, node):
511 self.write(left)
512 self.visit(node.elt)
513 for comprehension in node.generators:
514 self.visit(comprehension)
515 self.write(right)
516 return visit
517
518 visit_ListComp = generator_visit('[', ']')
519 visit_GeneratorExp = generator_visit('(', ')')
520 visit_SetComp = generator_visit('{', '}')
521 del generator_visit
522
523 def visit_DictComp(self, node):
524 self.write('{')
525 self.visit(node.key)
526 self.write(': ')
527 self.visit(node.value)
528 for comprehension in node.generators:
529 self.visit(comprehension)
530 self.write('}')
531
532 def visit_IfExp(self, node):
533 self.visit(node.body)
534 self.write(' if ')
535 self.visit(node.test)
536 self.write(' else ')
537 self.visit(node.orelse)
538
539 def visit_Starred(self, node):
540 self.write('*')
541 self.visit(node.value)
542
543 def visit_Repr(self, node):
544 # XXX: python 2.6 only
545 self.write('`')
546 self.visit(node.value)
547 self.write('`')
548
549 # Helper Nodes
550
551 def visit_alias(self, node):
552 self.write(node.name)
553 if node.asname is not None:
554 self.write(' as ' + node.asname)
555
556 def visit_comprehension(self, node):
557 self.write(' for ')
558 self.visit(node.target)
559 self.write(' in ')
560 self.visit(node.iter)
561 if node.ifs:
562 for if_ in node.ifs:
563 self.write(' if ')
564 self.visit(if_)
565
566 def visit_excepthandler(self, node):
567 self.newline(node)
568 self.write('except')
569 if node.type is not None:
570 self.write(' ')
571 self.visit(node.type)
572 if node.name is not None:
573 self.write(' as ')
574 self.visit(node.name)
575 self.write(':')
576 self.body(node.body)