Simple example of how to use ast.NodeVisitor?

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.

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()

generic_visit is called when a custom visitor (ie visit_Name) can't be found. Here's a piece of code I wrote recently with ast.NodeVisitor: https://foss.heptapod.net/pypy/pypy/-/blob/80ead76ab428100ffeb01109c7fc0d94f1048af2/py/_code/_assertionnew.py It interprets the AST nodes to gain debugging information about some of them and falls back in with generic_visit when a special implementation isn't provided.

Simple example of how to use ast.NodeVisitor?

Related

Recent Posts