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