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