| 1 | 
 | 
| 2 | # XXX should probably rename ASTVisitor to ASTWalker
 | 
| 3 | # XXX can it be made even more generic?
 | 
| 4 | 
 | 
| 5 | class ASTVisitor(object):
 | 
| 6 |     """Performs a depth-first walk of the AST
 | 
| 7 | 
 | 
| 8 |     The ASTVisitor will walk the AST, performing either a preorder or
 | 
| 9 |     postorder traversal depending on which method is called.
 | 
| 10 | 
 | 
| 11 |     methods:
 | 
| 12 |     preorder(tree, visitor)
 | 
| 13 |     postorder(tree, visitor)
 | 
| 14 |         tree: an instance of ast.Node
 | 
| 15 |         visitor: an instance with visitXXX methods
 | 
| 16 | 
 | 
| 17 |     The ASTVisitor is responsible for walking over the tree in the
 | 
| 18 |     correct order.  For each node, it checks the visitor argument for
 | 
| 19 |     a method named 'visitNodeType' where NodeType is the name of the
 | 
| 20 |     node's class, e.g. Class.  If the method exists, it is called
 | 
| 21 |     with the node as its sole argument.
 | 
| 22 | 
 | 
| 23 |     The visitor method for a particular node type can control how
 | 
| 24 |     child nodes are visited during a preorder walk.  (It can't control
 | 
| 25 |     the order during a postorder walk, because it is called _after_
 | 
| 26 |     the walk has occurred.)  The ASTVisitor modifies the visitor
 | 
| 27 |     argument by adding a visit method to the visitor; this method can
 | 
| 28 |     be used to visit a child node of arbitrary type.
 | 
| 29 |     """
 | 
| 30 | 
 | 
| 31 |     VERBOSE = 0
 | 
| 32 | 
 | 
| 33 |     def __init__(self):
 | 
| 34 |         self._method_cache = {}
 | 
| 35 | 
 | 
| 36 |     def _Default(self, node, *args):
 | 
| 37 |         """If a visitClassName method isn't provided, visit children."""
 | 
| 38 |         for child in node.getChildNodes():
 | 
| 39 |             self.Dispatch(child, *args)
 | 
| 40 | 
 | 
| 41 |     def Dispatch(self, node, *args):
 | 
| 42 |         klass = node.__class__
 | 
| 43 | 
 | 
| 44 |         # TODO: Shouldn't it be keyed by string rather than class instance?
 | 
| 45 |         # This would probably change the bytecode order.
 | 
| 46 |         meth = self._method_cache.get(klass, None)
 | 
| 47 |         if meth is None:
 | 
| 48 |             className = klass.__name__
 | 
| 49 |             meth = getattr(self, 'visit' + className, self._Default)
 | 
| 50 |             self._method_cache[klass] = meth
 | 
| 51 |         return meth(node, *args)
 | 
| 52 | 
 | 
| 53 |     # Subclasses call self.visit().  TODO: Rename?
 | 
| 54 |     visit = Dispatch
 | 
| 55 | 
 | 
| 56 | 
 | 
| 57 | class ExampleASTVisitor(ASTVisitor):
 | 
| 58 |     """Prints examples of the nodes that aren't visited
 | 
| 59 | 
 | 
| 60 |     This visitor-driver is only useful for development, when it's
 | 
| 61 |     helpful to develop a visitor incrementally, and get feedback on what
 | 
| 62 |     you still have to do.
 | 
| 63 |     """
 | 
| 64 |     examples = {}
 | 
| 65 | 
 | 
| 66 |     def Dispatch(self, node, *args):
 | 
| 67 |         self.node = node
 | 
| 68 |         meth = self._method_cache.get(node.__class__, None)
 | 
| 69 |         className = node.__class__.__name__
 | 
| 70 |         if meth is None:
 | 
| 71 |             meth = getattr(self.visitor, 'visit' + className, 0)
 | 
| 72 |             self._method_cache[node.__class__] = meth
 | 
| 73 |         if self.VERBOSE > 1:
 | 
| 74 |             print("Dispatch", className, (meth and meth.__name__ or ''))
 | 
| 75 |         if meth:
 | 
| 76 |             meth(node, *args)
 | 
| 77 |         elif self.VERBOSE > 0:
 | 
| 78 |             klass = node.__class__
 | 
| 79 |             if klass not in self.examples:
 | 
| 80 |                 self.examples[klass] = klass
 | 
| 81 |                 print()
 | 
| 82 |                 print(self.visitor)
 | 
| 83 |                 print(klass)
 | 
| 84 |                 for attr in dir(node):
 | 
| 85 |                     if attr[0] != '_':
 | 
| 86 |                         print("\t", "%-12.12s" % attr, getattr(node, attr))
 | 
| 87 |                 print()
 | 
| 88 |             return self._Default(node, *args)
 |