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