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