1# -*- coding: utf-8 -*- 2# 3# Copyright (C) 2008-2009 Edgewall Software 4# All rights reserved. 5# 6# This software is licensed as described in the file COPYING, which 7# you should have received as part of this distribution. The terms 8# are also available at http://genshi.edgewall.org/wiki/License. 9# 10# This software consists of voluntary contributions made by many 11# individuals. For the exact contribution history, see the revision 12# history and logs, available at http://genshi.edgewall.org/log/. 13 14"""Support classes for generating code from abstract syntax trees.""" 15 16try: 17 import ast 18except ImportError: 19 from chameleon import ast25 as ast 20 21import sys 22import logging 23import weakref 24import collections 25 26node_annotations = weakref.WeakKeyDictionary() 27 28try: 29 node_annotations[ast.Name()] = None 30except TypeError: 31 logging.debug( 32 "Unable to create weak references to AST nodes. " \ 33 "A lock will be used around compilation loop." 34 ) 35 36 node_annotations = {} 37 38__docformat__ = 'restructuredtext en' 39 40 41def annotated(value): 42 node = load("annotation") 43 node_annotations[node] = value 44 return node 45 46 47def parse(source, mode='eval'): 48 return compile(source, '', mode, ast.PyCF_ONLY_AST) 49 50 51def load(name): 52 return ast.Name(id=name, ctx=ast.Load()) 53 54 55def store(name): 56 return ast.Name(id=name, ctx=ast.Store()) 57 58 59def param(name): 60 return ast.Name(id=name, ctx=ast.Param()) 61 62 63def delete(name): 64 return ast.Name(id=name, ctx=ast.Del()) 65 66 67def subscript(name, value, ctx): 68 return ast.Subscript( 69 value=value, 70 slice=ast.Index(value=ast.Str(s=name)), 71 ctx=ctx, 72 ) 73 74 75def walk_names(target, mode): 76 for node in ast.walk(target): 77 if isinstance(node, ast.Name) and \ 78 isinstance(node.ctx, mode): 79 yield node.id 80 81 82def iter_fields(node): 83 """ 84 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields`` 85 that is present on *node*. 86 """ 87 for field in node._fields: 88 try: 89 yield field, getattr(node, field) 90 except AttributeError: 91 pass 92 93 94def iter_child_nodes(node): 95 """ 96 Yield all direct child nodes of *node*, that is, all fields that are nodes 97 and all items of fields that are lists of nodes. 98 """ 99 for name, field in iter_fields(node): 100 if isinstance(field, Node): 101 yield field 102 elif isinstance(field, list): 103 for item in field: 104 if isinstance(item, Node): 105 yield item 106 107 108def walk(node): 109 """ 110 Recursively yield all descendant nodes in the tree starting at *node* 111 (including *node* itself), in no specified order. This is useful if you 112 only want to modify nodes in place and don't care about the context. 113 """ 114 todo = collections.deque([node]) 115 while todo: 116 node = todo.popleft() 117 todo.extend(iter_child_nodes(node)) 118 yield node 119 120 121def copy(source, target): 122 target.__class__ = source.__class__ 123 target.__dict__ = source.__dict__ 124 125 126def swap(root, replacement, name): 127 for node in ast.walk(root): 128 if (isinstance(node, ast.Name) and 129 isinstance(node.ctx, ast.Load) and 130 node.id == name): 131 assert hasattr(replacement, '_fields') 132 node_annotations.setdefault(node, replacement) 133 134 135def marker(name): 136 return ast.Str(s="__%s" % name) 137 138 139class Node(object): 140 """AST baseclass that gives us a convenient initialization 141 method. We explicitly declare and use the ``_fields`` attribute.""" 142 143 _fields = () 144 145 def __init__(self, *args, **kwargs): 146 assert isinstance(self._fields, tuple) 147 self.__dict__.update(kwargs) 148 for name, value in zip(self._fields, args): 149 setattr(self, name, value) 150 151 def __repr__(self): 152 """Poor man's single-line pretty printer.""" 153 154 name = type(self).__name__ 155 return '<%s%s at %x>' % ( 156 name, 157 "".join(" %s=%r" % (name, getattr(self, name, "\"?\"")) 158 for name in self._fields), 159 id(self) 160 ) 161 162 def extract(self, condition): 163 result = [] 164 for node in walk(self): 165 if condition(node): 166 result.append(node) 167 168 return result 169 170 171class Builtin(Node): 172 """Represents a Python builtin. 173 174 Used when a builtin is used internally by the compiler, to avoid 175 clashing with a user assignment (e.g. ``help`` is a builtin, but 176 also commonly assigned in templates). 177 """ 178 179 _fields = "id", "ctx" 180 181 ctx = ast.Load() 182 183 184class Symbol(Node): 185 """Represents an importable symbol.""" 186 187 _fields = "value", 188 189 190class Static(Node): 191 """Represents a static value.""" 192 193 _fields = "value", "name" 194 195 name = None 196 197 198class Comment(Node): 199 _fields = "text", "space", "stmt" 200 201 stmt = None 202 space = "" 203 204 205class TokenRef(Node): 206 """Represents a source-code token reference.""" 207 208 _fields = "pos", "length" 209 210 211class ASTCodeGenerator(object): 212 """General purpose base class for AST transformations. 213 214 Every visitor method can be overridden to return an AST node that has been 215 altered or replaced in some way. 216 """ 217 218 def __init__(self, tree): 219 self.lines_info = [] 220 self.line_info = [] 221 self.lines = [] 222 self.line = "" 223 self.last = None 224 self.indent = 0 225 self.blame_stack = [] 226 self.visit(tree) 227 228 if self.line.strip(): 229 self._new_line() 230 231 self.line = None 232 self.line_info = None 233 234 # strip trivial lines 235 self.code = "\n".join( 236 line.strip() and line or "" 237 for line in self.lines 238 ) 239 240 def _change_indent(self, delta): 241 self.indent += delta 242 243 def _new_line(self): 244 if self.line is not None: 245 self.lines.append(self.line) 246 self.lines_info.append(self.line_info) 247 self.line = ' ' * 4 * self.indent 248 if len(self.blame_stack) == 0: 249 self.line_info = [] 250 self.last = None 251 else: 252 self.line_info = [(0, self.blame_stack[-1],)] 253 self.last = self.blame_stack[-1] 254 255 def _write(self, s): 256 if len(s) == 0: 257 return 258 if len(self.blame_stack) == 0: 259 if self.last is not None: 260 self.last = None 261 self.line_info.append((len(self.line), self.last)) 262 else: 263 if self.last != self.blame_stack[-1]: 264 self.last = self.blame_stack[-1] 265 self.line_info.append((len(self.line), self.last)) 266 self.line += s 267 268 def flush(self): 269 if self.line: 270 self._new_line() 271 272 def visit(self, node): 273 if node is None: 274 return None 275 if type(node) is tuple: 276 return tuple([self.visit(n) for n in node]) 277 try: 278 self.blame_stack.append((node.lineno, node.col_offset,)) 279 info = True 280 except AttributeError: 281 info = False 282 visitor = getattr(self, 'visit_%s' % node.__class__.__name__, None) 283 if visitor is None: 284 raise Exception('No handler for ``%s`` (%s).' % ( 285 node.__class__.__name__, repr(node))) 286 ret = visitor(node) 287 if info: 288 self.blame_stack.pop() 289 return ret 290 291 def visit_Module(self, node): 292 for n in node.body: 293 self.visit(n) 294 visit_Interactive = visit_Module 295 visit_Suite = visit_Module 296 297 def visit_Expression(self, node): 298 return self.visit(node.body) 299 300 # arguments = (expr* args, identifier? vararg, 301 # identifier? kwarg, expr* defaults) 302 def visit_arguments(self, node): 303 first = True 304 no_default_count = len(node.args) - len(node.defaults) 305 for i, arg in enumerate(node.args): 306 if not first: 307 self._write(', ') 308 else: 309 first = False 310 self.visit(arg) 311 if i >= no_default_count: 312 self._write('=') 313 self.visit(node.defaults[i - no_default_count]) 314 if getattr(node, 'vararg', None): 315 if not first: 316 self._write(', ') 317 else: 318 first = False 319 self._write('*' + node.vararg) 320 if getattr(node, 'kwarg', None): 321 if not first: 322 self._write(', ') 323 else: 324 first = False 325 self._write('**' + node.kwarg) 326 327 def visit_arg(self, node): 328 self._write(node.arg) 329 330 # FunctionDef(identifier name, arguments args, 331 # stmt* body, expr* decorators) 332 def visit_FunctionDef(self, node): 333 self._new_line() 334 for decorator in getattr(node, 'decorator_list', ()): 335 self._new_line() 336 self._write('@') 337 self.visit(decorator) 338 self._new_line() 339 self._write('def ' + node.name + '(') 340 self.visit(node.args) 341 self._write('):') 342 self._change_indent(1) 343 for statement in node.body: 344 self.visit(statement) 345 self._change_indent(-1) 346 347 # ClassDef(identifier name, expr* bases, stmt* body) 348 def visit_ClassDef(self, node): 349 self._new_line() 350 self._write('class ' + node.name) 351 if node.bases: 352 self._write('(') 353 self.visit(node.bases[0]) 354 for base in node.bases[1:]: 355 self._write(', ') 356 self.visit(base) 357 self._write(')') 358 self._write(':') 359 self._change_indent(1) 360 for statement in node.body: 361 self.visit(statement) 362 self._change_indent(-1) 363 364 # Return(expr? value) 365 def visit_Return(self, node): 366 self._new_line() 367 self._write('return') 368 if getattr(node, 'value', None): 369 self._write(' ') 370 self.visit(node.value) 371 372 # Delete(expr* targets) 373 def visit_Delete(self, node): 374 self._new_line() 375 self._write('del ') 376 self.visit(node.targets[0]) 377 for target in node.targets[1:]: 378 self._write(', ') 379 self.visit(target) 380 381 # Assign(expr* targets, expr value) 382 def visit_Assign(self, node): 383 self._new_line() 384 for target in node.targets: 385 self.visit(target) 386 self._write(' = ') 387 self.visit(node.value) 388 389 # AugAssign(expr target, operator op, expr value) 390 def visit_AugAssign(self, node): 391 self._new_line() 392 self.visit(node.target) 393 self._write(' ' + self.binary_operators[node.op.__class__] + '= ') 394 self.visit(node.value) 395 396 # Print(expr? dest, expr* values, bool nl) 397 def visit_Print(self, node): 398 self._new_line() 399 self._write('print') 400 if getattr(node, 'dest', None): 401 self._write(' >> ') 402 self.visit(node.dest) 403 if getattr(node, 'values', None): 404 self._write(', ') 405 else: 406 self._write(' ') 407 if getattr(node, 'values', None): 408 self.visit(node.values[0]) 409 for value in node.values[1:]: 410 self._write(', ') 411 self.visit(value) 412 if not node.nl: 413 self._write(',') 414 415 # For(expr target, expr iter, stmt* body, stmt* orelse) 416 def visit_For(self, node): 417 self._new_line() 418 self._write('for ') 419 self.visit(node.target) 420 self._write(' in ') 421 self.visit(node.iter) 422 self._write(':') 423 self._change_indent(1) 424 for statement in node.body: 425 self.visit(statement) 426 self._change_indent(-1) 427 if getattr(node, 'orelse', None): 428 self._new_line() 429 self._write('else:') 430 self._change_indent(1) 431 for statement in node.orelse: 432 self.visit(statement) 433 self._change_indent(-1) 434 435 # While(expr test, stmt* body, stmt* orelse) 436 def visit_While(self, node): 437 self._new_line() 438 self._write('while ') 439 self.visit(node.test) 440 self._write(':') 441 self._change_indent(1) 442 for statement in node.body: 443 self.visit(statement) 444 self._change_indent(-1) 445 if getattr(node, 'orelse', None): 446 self._new_line() 447 self._write('else:') 448 self._change_indent(1) 449 for statement in node.orelse: 450 self.visit(statement) 451 self._change_indent(-1) 452 453 # If(expr test, stmt* body, stmt* orelse) 454 def visit_If(self, node): 455 self._new_line() 456 self._write('if ') 457 self.visit(node.test) 458 self._write(':') 459 self._change_indent(1) 460 for statement in node.body: 461 self.visit(statement) 462 self._change_indent(-1) 463 if getattr(node, 'orelse', None): 464 self._new_line() 465 self._write('else:') 466 self._change_indent(1) 467 for statement in node.orelse: 468 self.visit(statement) 469 self._change_indent(-1) 470 471 # With(expr context_expr, expr? optional_vars, stmt* body) 472 def visit_With(self, node): 473 self._new_line() 474 self._write('with ') 475 self.visit(node.context_expr) 476 if getattr(node, 'optional_vars', None): 477 self._write(' as ') 478 self.visit(node.optional_vars) 479 self._write(':') 480 self._change_indent(1) 481 for statement in node.body: 482 self.visit(statement) 483 self._change_indent(-1) 484 485 # Raise(expr? type, expr? inst, expr? tback) 486 def visit_Raise(self, node): 487 self._new_line() 488 self._write('raise') 489 if not getattr(node, "type", None): 490 exc = getattr(node, "exc", None) 491 if exc is None: 492 return 493 self._write(' ') 494 return self.visit(exc) 495 self._write(' ') 496 self.visit(node.type) 497 if not node.inst: 498 return 499 self._write(', ') 500 self.visit(node.inst) 501 if not node.tback: 502 return 503 self._write(', ') 504 self.visit(node.tback) 505 506 # Try(stmt* body, excepthandler* handlers, stmt* orelse, stmt* finalbody) 507 def visit_Try(self, node): 508 self._new_line() 509 self._write('try:') 510 self._change_indent(1) 511 for statement in node.body: 512 self.visit(statement) 513 self._change_indent(-1) 514 if getattr(node, 'handlers', None): 515 for handler in node.handlers: 516 self.visit(handler) 517 self._new_line() 518 519 if getattr(node, 'orelse', None): 520 self._write('else:') 521 self._change_indent(1) 522 for statement in node.orelse: 523 self.visit(statement) 524 self._change_indent(-1) 525 526 if getattr(node, 'finalbody', None): 527 self._new_line() 528 self._write('finally:') 529 self._change_indent(1) 530 for statement in node.finalbody: 531 self.visit(statement) 532 self._change_indent(-1) 533 534 # TryExcept(stmt* body, excepthandler* handlers, stmt* orelse) 535 def visit_TryExcept(self, node): 536 self._new_line() 537 self._write('try:') 538 self._change_indent(1) 539 for statement in node.body: 540 self.visit(statement) 541 self._change_indent(-1) 542 if getattr(node, 'handlers', None): 543 for handler in node.handlers: 544 self.visit(handler) 545 self._new_line() 546 if getattr(node, 'orelse', None): 547 self._write('else:') 548 self._change_indent(1) 549 for statement in node.orelse: 550 self.visit(statement) 551 self._change_indent(-1) 552 553 # excepthandler = (expr? type, expr? name, stmt* body) 554 def visit_ExceptHandler(self, node): 555 self._new_line() 556 self._write('except') 557 if getattr(node, 'type', None): 558 self._write(' ') 559 self.visit(node.type) 560 if getattr(node, 'name', None): 561 if sys.version_info[0] == 2: 562 assert getattr(node, 'type', None) 563 self._write(', ') 564 else: 565 self._write(' as ') 566 self.visit(node.name) 567 self._write(':') 568 self._change_indent(1) 569 for statement in node.body: 570 self.visit(statement) 571 self._change_indent(-1) 572 visit_excepthandler = visit_ExceptHandler 573 574 # TryFinally(stmt* body, stmt* finalbody) 575 def visit_TryFinally(self, node): 576 self._new_line() 577 self._write('try:') 578 self._change_indent(1) 579 for statement in node.body: 580 self.visit(statement) 581 self._change_indent(-1) 582 583 if getattr(node, 'finalbody', None): 584 self._new_line() 585 self._write('finally:') 586 self._change_indent(1) 587 for statement in node.finalbody: 588 self.visit(statement) 589 self._change_indent(-1) 590 591 # Assert(expr test, expr? msg) 592 def visit_Assert(self, node): 593 self._new_line() 594 self._write('assert ') 595 self.visit(node.test) 596 if getattr(node, 'msg', None): 597 self._write(', ') 598 self.visit(node.msg) 599 600 def visit_alias(self, node): 601 self._write(node.name) 602 if getattr(node, 'asname', None): 603 self._write(' as ') 604 self._write(node.asname) 605 606 # Import(alias* names) 607 def visit_Import(self, node): 608 self._new_line() 609 self._write('import ') 610 self.visit(node.names[0]) 611 for name in node.names[1:]: 612 self._write(', ') 613 self.visit(name) 614 615 # ImportFrom(identifier module, alias* names, int? level) 616 def visit_ImportFrom(self, node): 617 self._new_line() 618 self._write('from ') 619 if node.level: 620 self._write('.' * node.level) 621 self._write(node.module) 622 self._write(' import ') 623 self.visit(node.names[0]) 624 for name in node.names[1:]: 625 self._write(', ') 626 self.visit(name) 627 628 # Exec(expr body, expr? globals, expr? locals) 629 def visit_Exec(self, node): 630 self._new_line() 631 self._write('exec ') 632 self.visit(node.body) 633 if not node.globals: 634 return 635 self._write(', ') 636 self.visit(node.globals) 637 if not node.locals: 638 return 639 self._write(', ') 640 self.visit(node.locals) 641 642 # Global(identifier* names) 643 def visit_Global(self, node): 644 self._new_line() 645 self._write('global ') 646 self.visit(node.names[0]) 647 for name in node.names[1:]: 648 self._write(', ') 649 self.visit(name) 650 651 # Expr(expr value) 652 def visit_Expr(self, node): 653 self._new_line() 654 self.visit(node.value) 655 656 # Pass 657 def visit_Pass(self, node): 658 self._new_line() 659 self._write('pass') 660 661 # Break 662 def visit_Break(self, node): 663 self._new_line() 664 self._write('break') 665 666 # Continue 667 def visit_Continue(self, node): 668 self._new_line() 669 self._write('continue') 670 671 ### EXPRESSIONS 672 def with_parens(f): 673 def _f(self, node): 674 self._write('(') 675 f(self, node) 676 self._write(')') 677 return _f 678 679 bool_operators = {ast.And: 'and', ast.Or: 'or'} 680 681 # BoolOp(boolop op, expr* values) 682 @with_parens 683 def visit_BoolOp(self, node): 684 joiner = ' ' + self.bool_operators[node.op.__class__] + ' ' 685 self.visit(node.values[0]) 686 for value in node.values[1:]: 687 self._write(joiner) 688 self.visit(value) 689 690 binary_operators = { 691 ast.Add: '+', 692 ast.Sub: '-', 693 ast.Mult: '*', 694 ast.Div: '/', 695 ast.Mod: '%', 696 ast.Pow: '**', 697 ast.LShift: '<<', 698 ast.RShift: '>>', 699 ast.BitOr: '|', 700 ast.BitXor: '^', 701 ast.BitAnd: '&', 702 ast.FloorDiv: '//' 703 } 704 705 # BinOp(expr left, operator op, expr right) 706 @with_parens 707 def visit_BinOp(self, node): 708 self.visit(node.left) 709 self._write(' ' + self.binary_operators[node.op.__class__] + ' ') 710 self.visit(node.right) 711 712 unary_operators = { 713 ast.Invert: '~', 714 ast.Not: 'not', 715 ast.UAdd: '+', 716 ast.USub: '-', 717 } 718 719 # UnaryOp(unaryop op, expr operand) 720 def visit_UnaryOp(self, node): 721 self._write(self.unary_operators[node.op.__class__] + ' ') 722 self.visit(node.operand) 723 724 # Lambda(arguments args, expr body) 725 @with_parens 726 def visit_Lambda(self, node): 727 self._write('lambda ') 728 self.visit(node.args) 729 self._write(': ') 730 self.visit(node.body) 731 732 # IfExp(expr test, expr body, expr orelse) 733 @with_parens 734 def visit_IfExp(self, node): 735 self.visit(node.body) 736 self._write(' if ') 737 self.visit(node.test) 738 self._write(' else ') 739 self.visit(node.orelse) 740 741 # Dict(expr* keys, expr* values) 742 def visit_Dict(self, node): 743 self._write('{') 744 for key, value in zip(node.keys, node.values): 745 self.visit(key) 746 self._write(': ') 747 self.visit(value) 748 self._write(', ') 749 self._write('}') 750 751 def visit_Set(self, node): 752 self._write('{') 753 elts = list(node.elts) 754 last = elts.pop() 755 for elt in elts: 756 self.visit(elt) 757 self._write(', ') 758 self.visit(last) 759 self._write('}') 760 761 # ListComp(expr elt, comprehension* generators) 762 def visit_ListComp(self, node): 763 self._write('[') 764 self.visit(node.elt) 765 for generator in node.generators: 766 # comprehension = (expr target, expr iter, expr* ifs) 767 self._write(' for ') 768 self.visit(generator.target) 769 self._write(' in ') 770 self.visit(generator.iter) 771 for ifexpr in generator.ifs: 772 self._write(' if ') 773 self.visit(ifexpr) 774 self._write(']') 775 776 # GeneratorExp(expr elt, comprehension* generators) 777 def visit_GeneratorExp(self, node): 778 self._write('(') 779 self.visit(node.elt) 780 for generator in node.generators: 781 # comprehension = (expr target, expr iter, expr* ifs) 782 self._write(' for ') 783 self.visit(generator.target) 784 self._write(' in ') 785 self.visit(generator.iter) 786 for ifexpr in generator.ifs: 787 self._write(' if ') 788 self.visit(ifexpr) 789 self._write(')') 790 791 # Yield(expr? value) 792 def visit_Yield(self, node): 793 self._write('yield') 794 if getattr(node, 'value', None): 795 self._write(' ') 796 self.visit(node.value) 797 798 comparison_operators = { 799 ast.Eq: '==', 800 ast.NotEq: '!=', 801 ast.Lt: '<', 802 ast.LtE: '<=', 803 ast.Gt: '>', 804 ast.GtE: '>=', 805 ast.Is: 'is', 806 ast.IsNot: 'is not', 807 ast.In: 'in', 808 ast.NotIn: 'not in', 809 } 810 811 # Compare(expr left, cmpop* ops, expr* comparators) 812 @with_parens 813 def visit_Compare(self, node): 814 self.visit(node.left) 815 for op, comparator in zip(node.ops, node.comparators): 816 self._write(' ' + self.comparison_operators[op.__class__] + ' ') 817 self.visit(comparator) 818 819 # Call(expr func, expr* args, keyword* keywords, 820 # expr? starargs, expr? kwargs) 821 def visit_Call(self, node): 822 self.visit(node.func) 823 self._write('(') 824 first = True 825 for arg in node.args: 826 if not first: 827 self._write(', ') 828 first = False 829 self.visit(arg) 830 831 for keyword in node.keywords: 832 if not first: 833 self._write(', ') 834 first = False 835 # keyword = (identifier arg, expr value) 836 if keyword.arg is not None: 837 self._write(keyword.arg) 838 self._write('=') 839 else: 840 self._write('**') 841 self.visit(keyword.value) 842 # Attribute removed in Python 3.5 843 if getattr(node, 'starargs', None): 844 if not first: 845 self._write(', ') 846 first = False 847 self._write('*') 848 self.visit(node.starargs) 849 850 # Attribute removed in Python 3.5 851 if getattr(node, 'kwargs', None): 852 if not first: 853 self._write(', ') 854 first = False 855 self._write('**') 856 self.visit(node.kwargs) 857 self._write(')') 858 859 # Repr(expr value) 860 def visit_Repr(self, node): 861 self._write('`') 862 self.visit(node.value) 863 self._write('`') 864 865 # Constant(object value) 866 def visit_Constant(self, node): 867 if node.value is Ellipsis: 868 self._write('...') 869 else: 870 self._write(repr(node.value)) 871 872 # Num(object n) 873 def visit_Num(self, node): 874 self._write(repr(node.n)) 875 876 # Str(string s) 877 def visit_Str(self, node): 878 self._write(repr(node.s)) 879 880 def visit_Ellipsis(self, node): 881 self._write('...') 882 883 # Attribute(expr value, identifier attr, expr_context ctx) 884 def visit_Attribute(self, node): 885 self.visit(node.value) 886 self._write('.') 887 self._write(node.attr) 888 889 # Subscript(expr value, slice slice, expr_context ctx) 890 def visit_Subscript(self, node): 891 self.visit(node.value) 892 self._write('[') 893 if isinstance(node.slice, ast.Tuple) and node.slice.elts: 894 self.visit(node.slice.elts[0]) 895 if len(node.slice.elts) == 1: 896 self._write(', ') 897 else: 898 for dim in node.slice.elts[1:]: 899 self._write(', ') 900 self.visit(dim) 901 else: 902 self.visit(node.slice) 903 self._write(']') 904 905 # Slice(expr? lower, expr? upper, expr? step) 906 def visit_Slice(self, node): 907 if getattr(node, 'lower', None) is not None: 908 self.visit(node.lower) 909 self._write(':') 910 if getattr(node, 'upper', None) is not None: 911 self.visit(node.upper) 912 if getattr(node, 'step', None) is not None: 913 self._write(':') 914 self.visit(node.step) 915 916 # Index(expr value) 917 def visit_Index(self, node): 918 self.visit(node.value) 919 920 # ExtSlice(slice* dims) 921 def visit_ExtSlice(self, node): 922 self.visit(node.dims[0]) 923 if len(node.dims) == 1: 924 self._write(', ') 925 else: 926 for dim in node.dims[1:]: 927 self._write(', ') 928 self.visit(dim) 929 930 # Starred(expr value, expr_context ctx) 931 def visit_Starred(self, node): 932 self._write('*') 933 self.visit(node.value) 934 935 # Name(identifier id, expr_context ctx) 936 def visit_Name(self, node): 937 self._write(node.id) 938 939 # List(expr* elts, expr_context ctx) 940 def visit_List(self, node): 941 self._write('[') 942 for elt in node.elts: 943 self.visit(elt) 944 self._write(', ') 945 self._write(']') 946 947 # Tuple(expr *elts, expr_context ctx) 948 def visit_Tuple(self, node): 949 self._write('(') 950 for elt in node.elts: 951 self.visit(elt) 952 self._write(', ') 953 self._write(')') 954 955 # NameConstant(singleton value) 956 def visit_NameConstant(self, node): 957 self._write(str(node.value)) 958 959class AnnotationAwareVisitor(ast.NodeVisitor): 960 def visit(self, node): 961 annotation = node_annotations.get(node) 962 if annotation is not None: 963 assert hasattr(annotation, '_fields') 964 node = annotation 965 966 super(AnnotationAwareVisitor, self).visit(node) 967 968 def apply_transform(self, node): 969 if node not in node_annotations: 970 result = self.transform(node) 971 if result is not None and result is not node: 972 node_annotations[node] = result 973 974 975class NameLookupRewriteVisitor(AnnotationAwareVisitor): 976 def __init__(self, transform): 977 self.transform = transform 978 self.transformed = set() 979 self.scopes = [set()] 980 981 def __call__(self, node): 982 self.visit(node) 983 return self.transformed 984 985 def visit_arg(self, node): 986 scope = self.scopes[-1] 987 scope.add(node.arg) 988 989 def visit_Name(self, node): 990 scope = self.scopes[-1] 991 if isinstance(node.ctx, ast.Param): 992 scope.add(node.id) 993 elif node.id not in scope: 994 self.transformed.add(node.id) 995 self.apply_transform(node) 996 997 def visit_FunctionDef(self, node): 998 self.scopes[-1].add(node.name) 999 1000 def visit_alias(self, node): 1001 name = node.asname if node.asname is not None else node.name 1002 self.scopes[-1].add(name) 1003 1004 def visit_Lambda(self, node): 1005 self.scopes.append(set()) 1006 try: 1007 self.visit(node.args) 1008 self.visit(node.body) 1009 finally: 1010 self.scopes.pop() 1011 1012 1013class ItemLookupOnAttributeErrorVisitor(AnnotationAwareVisitor): 1014 def __init__(self, transform): 1015 self.transform = transform 1016 1017 def visit_Attribute(self, node): 1018 self.generic_visit(node) 1019 self.apply_transform(node) 1020