blob: 018b2831775d52306449947312b0068a23768f1a [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
Andrew Geisslereff27472021-10-29 15:35:00 -0500395 def visit_Constant(self, node):
Andrew Geisslereff27472021-10-29 15:35:00 -0500396 self.write(repr(node.value))
397
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500398 def visit_Tuple(self, node):
399 self.write('(')
400 idx = -1
401 for idx, item in enumerate(node.elts):
402 if idx:
403 self.write(', ')
404 self.visit(item)
405 self.write(idx and ')' or ',)')
406
407 def sequence_visit(left, right):
408 def visit(self, node):
409 self.write(left)
410 for idx, item in enumerate(node.elts):
411 if idx:
412 self.write(', ')
413 self.visit(item)
414 self.write(right)
415 return visit
416
417 visit_List = sequence_visit('[', ']')
418 visit_Set = sequence_visit('{', '}')
419 del sequence_visit
420
421 def visit_Dict(self, node):
422 self.write('{')
423 for idx, (key, value) in enumerate(zip(node.keys, node.values)):
424 if idx:
425 self.write(', ')
426 self.visit(key)
427 self.write(': ')
428 self.visit(value)
429 self.write('}')
430
431 def visit_BinOp(self, node):
432 self.visit(node.left)
433 self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
434 self.visit(node.right)
435
436 def visit_BoolOp(self, node):
437 self.write('(')
438 for idx, value in enumerate(node.values):
439 if idx:
440 self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
441 self.visit(value)
442 self.write(')')
443
444 def visit_Compare(self, node):
445 self.write('(')
446 self.write(node.left)
447 for op, right in zip(node.ops, node.comparators):
448 self.write(' %s %%' % CMPOP_SYMBOLS[type(op)])
449 self.visit(right)
450 self.write(')')
451
452 def visit_UnaryOp(self, node):
453 self.write('(')
454 op = UNARYOP_SYMBOLS[type(node.op)]
455 self.write(op)
456 if op == 'not':
457 self.write(' ')
458 self.visit(node.operand)
459 self.write(')')
460
461 def visit_Subscript(self, node):
462 self.visit(node.value)
463 self.write('[')
464 self.visit(node.slice)
465 self.write(']')
466
467 def visit_Slice(self, node):
468 if node.lower is not None:
469 self.visit(node.lower)
470 self.write(':')
471 if node.upper is not None:
472 self.visit(node.upper)
473 if node.step is not None:
474 self.write(':')
475 if not (isinstance(node.step, Name) and node.step.id == 'None'):
476 self.visit(node.step)
477
478 def visit_ExtSlice(self, node):
479 for idx, item in node.dims:
480 if idx:
481 self.write(', ')
482 self.visit(item)
483
484 def visit_Yield(self, node):
485 self.write('yield ')
486 self.visit(node.value)
487
488 def visit_Lambda(self, node):
489 self.write('lambda ')
490 self.signature(node.args)
491 self.write(': ')
492 self.visit(node.body)
493
494 def visit_Ellipsis(self, node):
495 self.write('Ellipsis')
496
497 def generator_visit(left, right):
498 def visit(self, node):
499 self.write(left)
500 self.visit(node.elt)
501 for comprehension in node.generators:
502 self.visit(comprehension)
503 self.write(right)
504 return visit
505
506 visit_ListComp = generator_visit('[', ']')
507 visit_GeneratorExp = generator_visit('(', ')')
508 visit_SetComp = generator_visit('{', '}')
509 del generator_visit
510
511 def visit_DictComp(self, node):
512 self.write('{')
513 self.visit(node.key)
514 self.write(': ')
515 self.visit(node.value)
516 for comprehension in node.generators:
517 self.visit(comprehension)
518 self.write('}')
519
520 def visit_IfExp(self, node):
521 self.visit(node.body)
522 self.write(' if ')
523 self.visit(node.test)
524 self.write(' else ')
525 self.visit(node.orelse)
526
527 def visit_Starred(self, node):
528 self.write('*')
529 self.visit(node.value)
530
531 def visit_Repr(self, node):
532 # XXX: python 2.6 only
533 self.write('`')
534 self.visit(node.value)
535 self.write('`')
536
537 # Helper Nodes
538
539 def visit_alias(self, node):
540 self.write(node.name)
541 if node.asname is not None:
542 self.write(' as ' + node.asname)
543
544 def visit_comprehension(self, node):
545 self.write(' for ')
546 self.visit(node.target)
547 self.write(' in ')
548 self.visit(node.iter)
549 if node.ifs:
550 for if_ in node.ifs:
551 self.write(' if ')
552 self.visit(if_)
553
554 def visit_excepthandler(self, node):
555 self.newline(node)
556 self.write('except')
557 if node.type is not None:
558 self.write(' ')
559 self.visit(node.type)
560 if node.name is not None:
561 self.write(' as ')
562 self.visit(node.name)
563 self.write(':')
564 self.body(node.body)