10c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi# -*- coding: utf-8 -*-
20c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi"""
30c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    ast
40c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    ~~~
50c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
60c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    The `ast` module helps Python applications to process trees of the Python
70c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    abstract syntax grammar.  The abstract syntax itself might change with
80c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    each Python release; this module helps to find out programmatically what
90c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    the current grammar looks like and allows modifications of it.
100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    An abstract syntax tree can be generated by passing `ast.PyCF_ONLY_AST` as
120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    a flag to the `compile()` builtin function or by using the `parse()`
130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    function from this module.  The result will be a tree of objects whose
140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    classes all inherit from `ast.AST`.
150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    A modified abstract syntax tree can be compiled into a Python code object
170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    using the built-in `compile()` function.
180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Additionally various helper functions are provided that make working with
200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    the trees simpler.  The main intention of the helper functions and this
210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    module in general is to provide an easy to use interface for libraries
220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    that work tightly with the python syntax (template engines for example).
230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    :copyright: Copyright 2008 by Armin Ronacher.
260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    :license: Python License.
270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi"""
280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom _ast import *
290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom _ast import __version__
300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef parse(source, filename='<unknown>', mode='exec'):
330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Parse the source into an AST node.
350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return compile(source, filename, mode, PyCF_ONLY_AST)
380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef literal_eval(node_or_string):
410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Safely evaluate an expression node or a string containing a Python
430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    expression.  The string or node provided may only consist of the following
440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Python literal structures: strings, numbers, tuples, lists, dicts, booleans,
450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    and None.
460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    _safe_names = {'None': None, 'True': True, 'False': False}
480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    if isinstance(node_or_string, basestring):
490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        node_or_string = parse(node_or_string, mode='eval')
500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    if isinstance(node_or_string, Expression):
510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        node_or_string = node_or_string.body
520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def _convert(node):
530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if isinstance(node, Str):
540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return node.s
550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, Num):
560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return node.n
570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, Tuple):
580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return tuple(map(_convert, node.elts))
590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, List):
600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return list(map(_convert, node.elts))
610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, Dict):
620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return dict((_convert(k), _convert(v)) for k, v
630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                        in zip(node.keys, node.values))
640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, Name):
650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if node.id in _safe_names:
660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                return _safe_names[node.id]
670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, BinOp) and \
680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi             isinstance(node.op, (Add, Sub)) and \
690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi             isinstance(node.right, Num) and \
700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi             isinstance(node.right.n, complex) and \
710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi             isinstance(node.left, Num) and \
720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi             isinstance(node.left.n, (int, long, float)):
730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            left = node.left.n
740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            right = node.right.n
750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if isinstance(node.op, Add):
760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                return left + right
770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            else:
780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                return left - right
790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        raise ValueError('malformed string')
800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return _convert(node_or_string)
810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef dump(node, annotate_fields=True, include_attributes=False):
840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Return a formatted dump of the tree in *node*.  This is mainly useful for
860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    debugging purposes.  The returned string will show the names and the values
870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    for fields.  This makes the code impossible to evaluate, so if evaluation is
880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    wanted *annotate_fields* must be set to False.  Attributes such as line
890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    numbers and column offsets are not dumped by default.  If this is wanted,
900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    *include_attributes* can be set to True.
910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def _format(node):
930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if isinstance(node, AST):
940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            fields = [(a, _format(b)) for a, b in iter_fields(node)]
950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            rv = '%s(%s' % (node.__class__.__name__, ', '.join(
960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                ('%s=%s' % field for field in fields)
970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                if annotate_fields else
980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                (b for a, b in fields)
990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            ))
1000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if include_attributes and node._attributes:
1010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                rv += fields and ', ' or ' '
1020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                rv += ', '.join('%s=%s' % (a, _format(getattr(node, a)))
1030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                                for a in node._attributes)
1040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return rv + ')'
1050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(node, list):
1060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return '[%s]' % ', '.join(_format(x) for x in node)
1070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        return repr(node)
1080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    if not isinstance(node, AST):
1090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        raise TypeError('expected AST, got %r' % node.__class__.__name__)
1100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return _format(node)
1110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef copy_location(new_node, old_node):
1140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Copy source location (`lineno` and `col_offset` attributes) from
1160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    *old_node* to *new_node* if possible, and return *new_node*.
1170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    for attr in 'lineno', 'col_offset':
1190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if attr in old_node._attributes and attr in new_node._attributes \
1200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi           and hasattr(old_node, attr):
1210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            setattr(new_node, attr, getattr(old_node, attr))
1220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return new_node
1230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef fix_missing_locations(node):
1260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    When you compile a node tree with compile(), the compiler expects lineno and
1280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    col_offset attributes for every node that supports them.  This is rather
1290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    tedious to fill in for generated nodes, so this helper adds these attributes
1300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    recursively where not already set, by setting them to the values of the
1310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    parent node.  It works recursively starting at *node*.
1320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def _fix(node, lineno, col_offset):
1340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if 'lineno' in node._attributes:
1350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if not hasattr(node, 'lineno'):
1360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                node.lineno = lineno
1370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            else:
1380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                lineno = node.lineno
1390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if 'col_offset' in node._attributes:
1400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if not hasattr(node, 'col_offset'):
1410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                node.col_offset = col_offset
1420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            else:
1430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                col_offset = node.col_offset
1440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        for child in iter_child_nodes(node):
1450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            _fix(child, lineno, col_offset)
1460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    _fix(node, 1, 0)
1470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return node
1480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef increment_lineno(node, n=1):
1510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Increment the line number of each node in the tree starting at *node* by *n*.
1530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    This is useful to "move code" to a different location in a file.
1540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    for child in walk(node):
1560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if 'lineno' in child._attributes:
1570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            child.lineno = getattr(child, 'lineno', 0) + n
1580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return node
1590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef iter_fields(node):
1620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
1640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    that is present on *node*.
1650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    for field in node._fields:
1670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        try:
1680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            yield field, getattr(node, field)
1690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        except AttributeError:
1700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            pass
1710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef iter_child_nodes(node):
1740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Yield all direct child nodes of *node*, that is, all fields that are nodes
1760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    and all items of fields that are lists of nodes.
1770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    for name, field in iter_fields(node):
1790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if isinstance(field, AST):
1800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            yield field
1810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        elif isinstance(field, list):
1820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            for item in field:
1830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                if isinstance(item, AST):
1840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    yield item
1850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef get_docstring(node, clean=True):
1880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Return the docstring for the given node or None if no docstring can
1900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    be found.  If the node provided does not have docstrings a TypeError
1910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    will be raised.
1920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
1930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    if not isinstance(node, (FunctionDef, ClassDef, Module)):
1940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        raise TypeError("%r can't have docstrings" % node.__class__.__name__)
1950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    if node.body and isinstance(node.body[0], Expr) and \
1960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi       isinstance(node.body[0].value, Str):
1970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        if clean:
1980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            import inspect
1990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            return inspect.cleandoc(node.body[0].value.s)
2000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        return node.body[0].value.s
2010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef walk(node):
2040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Recursively yield all descendant nodes in the tree starting at *node*
2060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    (including *node* itself), in no specified order.  This is useful if you
2070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    only want to modify nodes in place and don't care about the context.
2080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    from collections import deque
2100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    todo = deque([node])
2110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    while todo:
2120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        node = todo.popleft()
2130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        todo.extend(iter_child_nodes(node))
2140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        yield node
2150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass NodeVisitor(object):
2180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    A node visitor base class that walks the abstract syntax tree and calls a
2200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    visitor function for every node found.  This function may return a value
2210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    which is forwarded by the `visit` method.
2220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    This class is meant to be subclassed, with the subclass adding visitor
2240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    methods.
2250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Per default the visitor functions for the nodes are ``'visit_'`` +
2270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    class name of the node.  So a `TryFinally` node visit function would
2280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    be `visit_TryFinally`.  This behavior can be changed by overriding
2290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    the `visit` method.  If no visitor function exists for a node
2300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    (return value `None`) the `generic_visit` visitor is used instead.
2310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Don't use the `NodeVisitor` if you want to apply changes to nodes during
2330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    traversing.  For this a special visitor exists (`NodeTransformer`) that
2340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    allows modifications.
2350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def visit(self, node):
2380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        """Visit a node."""
2390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        method = 'visit_' + node.__class__.__name__
2400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        visitor = getattr(self, method, self.generic_visit)
2410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        return visitor(node)
2420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def generic_visit(self, node):
2440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        """Called if no explicit visitor function exists for a node."""
2450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        for field, value in iter_fields(node):
2460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if isinstance(value, list):
2470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                for item in value:
2480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    if isinstance(item, AST):
2490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                        self.visit(item)
2500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            elif isinstance(value, AST):
2510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                self.visit(value)
2520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass NodeTransformer(NodeVisitor):
2550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    A :class:`NodeVisitor` subclass that walks the abstract syntax tree and
2570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    allows modification of nodes.
2580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    The `NodeTransformer` will walk the AST and use the return value of the
2600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    visitor methods to replace or remove the old node.  If the return value of
2610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    the visitor method is ``None``, the node will be removed from its location,
2620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    otherwise it is replaced with the return value.  The return value may be the
2630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    original node in which case no replacement takes place.
2640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Here is an example transformer that rewrites all occurrences of name lookups
2660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    (``foo``) to ``data['foo']``::
2670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi       class RewriteName(NodeTransformer):
2690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi           def visit_Name(self, node):
2710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi               return copy_location(Subscript(
2720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                   value=Name(id='data', ctx=Load()),
2730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                   slice=Index(value=Str(s=node.id)),
2740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                   ctx=node.ctx
2750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi               ), node)
2760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Keep in mind that if the node you're operating on has child nodes you must
2780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    either transform the child nodes yourself or call the :meth:`generic_visit`
2790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    method for the node first.
2800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    For nodes that were part of a collection of statements (that applies to all
2820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    statement nodes), the visitor may also return a list of nodes rather than
2830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    just a single node.
2840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    Usually you use the transformer like this::
2860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi       node = YourTransformer().visit(node)
2880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    """
2890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
2900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def generic_visit(self, node):
2910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        for field, old_value in iter_fields(node):
2920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            old_value = getattr(node, field, None)
2930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            if isinstance(old_value, list):
2940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                new_values = []
2950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                for value in old_value:
2960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    if isinstance(value, AST):
2970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                        value = self.visit(value)
2980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                        if value is None:
2990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                            continue
3000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                        elif not isinstance(value, AST):
3010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                            new_values.extend(value)
3020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                            continue
3030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    new_values.append(value)
3040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                old_value[:] = new_values
3050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            elif isinstance(old_value, AST):
3060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                new_node = self.visit(old_value)
3070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                if new_node is None:
3080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    delattr(node, field)
3090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                else:
3100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                    setattr(node, field, new_node)
3110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        return node
312