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