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