1edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep# -*- coding: utf-8 -*-
2edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep"""
3edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    ast
4edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    ~~~
5edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
6edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    The `ast` module helps Python applications to process trees of the Python
7edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    abstract syntax grammar.  The abstract syntax itself might change with
8edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    each Python release; this module helps to find out programmatically what
9edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    the current grammar looks like and allows modifications of it.
10edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
11edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    An abstract syntax tree can be generated by passing `ast.PyCF_ONLY_AST` as
12edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    a flag to the `compile()` builtin function or by using the `parse()`
13edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    function from this module.  The result will be a tree of objects whose
14edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    classes all inherit from `ast.AST`.
15edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
16edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    A modified abstract syntax tree can be compiled into a Python code object
17edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    using the built-in `compile()` function.
18edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
19edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Additionally various helper functions are provided that make working with
20edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    the trees simpler.  The main intention of the helper functions and this
21edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    module in general is to provide an easy to use interface for libraries
22edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    that work tightly with the python syntax (template engines for example).
23edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
24edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
25edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    :copyright: Copyright 2008 by Armin Ronacher.
26edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    :license: Python License.
27edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep"""
28edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepfrom _ast import *
29edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepfrom _ast import __version__
30edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
31edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
32edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef parse(source, filename='<unknown>', mode='exec'):
33edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
34edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Parse the source into an AST node.
35edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
36edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
37edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return compile(source, filename, mode, PyCF_ONLY_AST)
38edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
39edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
40edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef literal_eval(node_or_string):
41edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
42edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Safely evaluate an expression node or a string containing a Python
43edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    expression.  The string or node provided may only consist of the following
44edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Python literal structures: strings, numbers, tuples, lists, dicts, booleans,
45edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    and None.
46edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
47edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    _safe_names = {'None': None, 'True': True, 'False': False}
48edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    if isinstance(node_or_string, basestring):
49edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        node_or_string = parse(node_or_string, mode='eval')
50edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    if isinstance(node_or_string, Expression):
51edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        node_or_string = node_or_string.body
52edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def _convert(node):
53edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if isinstance(node, Str):
54edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return node.s
55edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, Num):
56edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return node.n
57edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, Tuple):
58edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return tuple(map(_convert, node.elts))
59edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, List):
60edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return list(map(_convert, node.elts))
61edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, Dict):
62edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return dict((_convert(k), _convert(v)) for k, v
63edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                        in zip(node.keys, node.values))
64edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, Name):
65edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if node.id in _safe_names:
66edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                return _safe_names[node.id]
67edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, BinOp) and \
68edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep             isinstance(node.op, (Add, Sub)) and \
69edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep             isinstance(node.right, Num) and \
70edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep             isinstance(node.right.n, complex) and \
71edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep             isinstance(node.left, Num) and \
72edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep             isinstance(node.left.n, (int, long, float)):
73edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            left = node.left.n
74edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            right = node.right.n
75edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if isinstance(node.op, Add):
76edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                return left + right
77edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            else:
78edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                return left - right
79edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        raise ValueError('malformed string')
80edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return _convert(node_or_string)
81edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
82edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
83edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef dump(node, annotate_fields=True, include_attributes=False):
84edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
85edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Return a formatted dump of the tree in *node*.  This is mainly useful for
86edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    debugging purposes.  The returned string will show the names and the values
87edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    for fields.  This makes the code impossible to evaluate, so if evaluation is
88edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    wanted *annotate_fields* must be set to False.  Attributes such as line
89edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    numbers and column offsets are not dumped by default.  If this is wanted,
90edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    *include_attributes* can be set to True.
91edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
92edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def _format(node):
93edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if isinstance(node, AST):
94edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            fields = [(a, _format(b)) for a, b in iter_fields(node)]
95edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            rv = '%s(%s' % (node.__class__.__name__, ', '.join(
96edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                ('%s=%s' % field for field in fields)
97edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                if annotate_fields else
98edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                (b for a, b in fields)
99edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            ))
100edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if include_attributes and node._attributes:
101edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                rv += fields and ', ' or ' '
102edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                rv += ', '.join('%s=%s' % (a, _format(getattr(node, a)))
103edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                                for a in node._attributes)
104edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return rv + ')'
105edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(node, list):
106edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return '[%s]' % ', '.join(_format(x) for x in node)
107edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        return repr(node)
108edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    if not isinstance(node, AST):
109edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        raise TypeError('expected AST, got %r' % node.__class__.__name__)
110edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return _format(node)
111edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
112edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
113edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef copy_location(new_node, old_node):
114edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
115edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Copy source location (`lineno` and `col_offset` attributes) from
116edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    *old_node* to *new_node* if possible, and return *new_node*.
117edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
118edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    for attr in 'lineno', 'col_offset':
119edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if attr in old_node._attributes and attr in new_node._attributes \
120edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep           and hasattr(old_node, attr):
121edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            setattr(new_node, attr, getattr(old_node, attr))
122edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return new_node
123edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
124edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
125edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef fix_missing_locations(node):
126edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
127edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    When you compile a node tree with compile(), the compiler expects lineno and
128edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    col_offset attributes for every node that supports them.  This is rather
129edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    tedious to fill in for generated nodes, so this helper adds these attributes
130edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    recursively where not already set, by setting them to the values of the
131edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    parent node.  It works recursively starting at *node*.
132edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
133edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def _fix(node, lineno, col_offset):
134edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if 'lineno' in node._attributes:
135edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if not hasattr(node, 'lineno'):
136edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                node.lineno = lineno
137edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            else:
138edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                lineno = node.lineno
139edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if 'col_offset' in node._attributes:
140edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if not hasattr(node, 'col_offset'):
141edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                node.col_offset = col_offset
142edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            else:
143edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                col_offset = node.col_offset
144edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        for child in iter_child_nodes(node):
145edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            _fix(child, lineno, col_offset)
146edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    _fix(node, 1, 0)
147edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return node
148edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
149edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
150edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef increment_lineno(node, n=1):
151edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
152edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Increment the line number of each node in the tree starting at *node* by *n*.
153edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    This is useful to "move code" to a different location in a file.
154edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
155edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    for child in walk(node):
156edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if 'lineno' in child._attributes:
157edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            child.lineno = getattr(child, 'lineno', 0) + n
158edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    return node
159edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
160edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
161edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef iter_fields(node):
162edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
163edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
164edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    that is present on *node*.
165edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
166edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    for field in node._fields:
167edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        try:
168edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            yield field, getattr(node, field)
169edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        except AttributeError:
170edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            pass
171edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
172edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
173edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef iter_child_nodes(node):
174edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
175edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Yield all direct child nodes of *node*, that is, all fields that are nodes
176edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    and all items of fields that are lists of nodes.
177edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
178edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    for name, field in iter_fields(node):
179edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if isinstance(field, AST):
180edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            yield field
181edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        elif isinstance(field, list):
182edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            for item in field:
183edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                if isinstance(item, AST):
184edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    yield item
185edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
186edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
187edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef get_docstring(node, clean=True):
188edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
189edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Return the docstring for the given node or None if no docstring can
190edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    be found.  If the node provided does not have docstrings a TypeError
191edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    will be raised.
192edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
193edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    if not isinstance(node, (FunctionDef, ClassDef, Module)):
194edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        raise TypeError("%r can't have docstrings" % node.__class__.__name__)
195edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    if node.body and isinstance(node.body[0], Expr) and \
196edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep       isinstance(node.body[0].value, Str):
197edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        if clean:
198edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            import inspect
199edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            return inspect.cleandoc(node.body[0].value.s)
200edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        return node.body[0].value.s
201edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
202edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
203edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef walk(node):
204edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
205edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Recursively yield all descendant nodes in the tree starting at *node*
206edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    (including *node* itself), in no specified order.  This is useful if you
207edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    only want to modify nodes in place and don't care about the context.
208edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
209edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    from collections import deque
210edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    todo = deque([node])
211edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    while todo:
212edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        node = todo.popleft()
213edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        todo.extend(iter_child_nodes(node))
214edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        yield node
215edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
216edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
217edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass NodeVisitor(object):
218edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
219edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    A node visitor base class that walks the abstract syntax tree and calls a
220edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    visitor function for every node found.  This function may return a value
221edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    which is forwarded by the `visit` method.
222edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
223edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    This class is meant to be subclassed, with the subclass adding visitor
224edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    methods.
225edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
226edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Per default the visitor functions for the nodes are ``'visit_'`` +
227edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    class name of the node.  So a `TryFinally` node visit function would
228edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    be `visit_TryFinally`.  This behavior can be changed by overriding
229edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    the `visit` method.  If no visitor function exists for a node
230edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    (return value `None`) the `generic_visit` visitor is used instead.
231edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
232edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Don't use the `NodeVisitor` if you want to apply changes to nodes during
233edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    traversing.  For this a special visitor exists (`NodeTransformer`) that
234edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    allows modifications.
235edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
236edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
237edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def visit(self, node):
238edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        """Visit a node."""
239edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        method = 'visit_' + node.__class__.__name__
240edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        visitor = getattr(self, method, self.generic_visit)
241edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        return visitor(node)
242edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
243edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def generic_visit(self, node):
244edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        """Called if no explicit visitor function exists for a node."""
245edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        for field, value in iter_fields(node):
246edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if isinstance(value, list):
247edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                for item in value:
248edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    if isinstance(item, AST):
249edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                        self.visit(item)
250edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            elif isinstance(value, AST):
251edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                self.visit(value)
252edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
253edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
254edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass NodeTransformer(NodeVisitor):
255edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
256edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
257edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    allows modification of nodes.
258edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
259edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    The `NodeTransformer` will walk the AST and use the return value of the
260edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    visitor methods to replace or remove the old node.  If the return value of
261edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    the visitor method is ``None``, the node will be removed from its location,
262edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    otherwise it is replaced with the return value.  The return value may be the
263edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    original node in which case no replacement takes place.
264edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
265edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Here is an example transformer that rewrites all occurrences of name lookups
266edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    (``foo``) to ``data['foo']``::
267edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
268edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep       class RewriteName(NodeTransformer):
269edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
270edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep           def visit_Name(self, node):
271edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep               return copy_location(Subscript(
272edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                   value=Name(id='data', ctx=Load()),
273edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                   slice=Index(value=Str(s=node.id)),
274edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                   ctx=node.ctx
275edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep               ), node)
276edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
277edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Keep in mind that if the node you're operating on has child nodes you must
278edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    either transform the child nodes yourself or call the :meth:`generic_visit`
279edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    method for the node first.
280edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
281edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    For nodes that were part of a collection of statements (that applies to all
282edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    statement nodes), the visitor may also return a list of nodes rather than
283edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    just a single node.
284edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
285edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    Usually you use the transformer like this::
286edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
287edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep       node = YourTransformer().visit(node)
288edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    """
289edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
290edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def generic_visit(self, node):
291edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        for field, old_value in iter_fields(node):
292edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            old_value = getattr(node, field, None)
293edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if isinstance(old_value, list):
294edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                new_values = []
295edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                for value in old_value:
296edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    if isinstance(value, AST):
297edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                        value = self.visit(value)
298edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                        if value is None:
299edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                            continue
300edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                        elif not isinstance(value, AST):
301edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                            new_values.extend(value)
302edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                            continue
303edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    new_values.append(value)
304edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                old_value[:] = new_values
305edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            elif isinstance(old_value, AST):
306edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                new_node = self.visit(old_value)
307edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                if new_node is None:
308edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    delattr(node, field)
309edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                else:
310edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                    setattr(node, field, new_node)
311edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        return node
312