1# -*- 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 hasattr(node, 'starargs') and node.starargs is not None: 218 paren_or_comma() 219 self.write('*') 220 self.visit(node.starargs) 221 if hasattr(node, 'kwargs') and 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 hasattr(node, 'starargs') and node.starargs is not None: 383 write_comma() 384 self.write('*') 385 self.visit(node.starargs) 386 if hasattr(node, 'kwargs') and 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_Constant(self, node): 396 self.write(repr(node.value)) 397 398 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) 565