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