python - 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.