Simple example of how to use ast.NodeVisitor?

Question

Does anyone have a simple example using ast.NodeVisitor to walk the abstract syntax tree in Python 2.6? The difference between visit and generic_visit is unclear to me, and I cannot find any example using google codesearch or plain google.

Answer 1

ast.visit -- unless you override it in a subclass, of course -- when called to visit an ast.Node of class foo, calls self.visit_foo if that method exists, otherwise self.generic_visit. The latter, again in its implementation in class ast itself, just calls self.visit on every child node (and performs no other action).

So, consider, for example:

>>> class v(ast.NodeVisitor): ...   def generic_visit(self, node): ...     print type(node).__name__ ...     ast.NodeVisitor.generic_visit(self, node) ...  

Here, we're overriding generic_visit to print the class name, but also calling up to the base class (so that all children will also be visited). So for example...:

>>> x = v() >>> t = ast.parse('d[x] += v[y, x]') >>> x.visit(t) 

emits:

Module AugAssign Subscript Name Load Index Name Load Store Add Subscript Name Load Index Tuple Name Load Name Load Load Load 

But suppose we didn't care for Load nodes (and children thereof -- if they had any;-). Then a simple way to deal with that might be, e.g.:

>>> class w(v): ...   def visit_Load(self, node): pass ...  

Now when we're visiting a Load node, visit dispatches, NOT to generic_visit any more, but to our new visit_Load... which doesn't do anything at all. So:

>>> y = w() >>> y.visit(t) Module AugAssign Subscript Name Index Name Store Add Subscript Name Index Tuple Name Name 

or, suppose we also wanted to see the actual names for Name nodes; then...:

>>> class z(v): ...   def visit_Name(self, node): print 'Name:', node.id ...  >>> z().visit(t) Module AugAssign Subscript Name: d Index Name: x Store Add Subscript Name: v Index Tuple Name: y Name: x Load Load 

But, NodeVisitor is a class because this lets it store information during a visit. Suppose all we want is the set of names in a "module". Then we don't need to override generic_visit any more, but rather...:

>>> class allnames(ast.NodeVisitor): ...   def visit_Module(self, node): ...     self.names = set() ...     self.generic_visit(node) ...     print sorted(self.names) ...   def visit_Name(self, node): ...     self.names.add(node.id) ...  >>> allnames().visit(t) ['d', 'v', 'x', 'y'] 

This kind of thing is a more typical use case than ones requiring overrides of generic_visit -- normally, you're only interested in a few kinds of nodes, like we are here in Module and Name, so we can just override visit_Module and visit_Name and let ast's visit do the dispatching on our behalf.

Answer 2

Looking at the code in ast.py it's not that hard to copy paste and roll your own walker. E.g.

import ast def str_node(node):     if isinstance(node, ast.AST):         fields = [(name, str_node(val)) for name, val in ast.iter_fields(node) if name not in ('left', 'right')]         rv = '%s(%s' % (node.__class__.__name__, ', '.join('%s=%s' % field for field in fields))         return rv + ')'     else:         return repr(node) def ast_visit(node, level=0):     print('  ' * level + str_node(node))     for field, value in ast.iter_fields(node):         if isinstance(value, list):             for item in value:                 if isinstance(item, ast.AST):                     ast_visit(item, level=level+1)         elif isinstance(value, ast.AST):             ast_visit(value, level=level+1)   ast_visit(ast.parse('a + b')) 

Prints out

Module(body=[<_ast.Expr object at 0x02808510>])   Expr(value=BinOp(op=Add()))     BinOp(op=Add())       Name(id='a', ctx=Load())         Load()       Add()       Name(id='b', ctx=Load())         Load()