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