1 from compiler import ast 2 3 # XXX should probably rename ASTVisitor to ASTWalker 4 # XXX can it be made even more generic? 5 6 class ASTVisitor: 7 """Performs a depth-first walk of the AST 8 9 The ASTVisitor will walk the AST, performing either a preorder or 10 postorder traversal depending on which method is called. 11 12 methods: 13 preorder(tree, visitor) 14 postorder(tree, visitor) 15 tree: an instance of ast.Node 16 visitor: an instance with visitXXX methods 17 18 The ASTVisitor is responsible for walking over the tree in the 19 correct order. For each node, it checks the visitor argument for 20 a method named 'visitNodeType' where NodeType is the name of the 21 node's class, e.g. Class. If the method exists, it is called 22 with the node as its sole argument. 23 24 The visitor method for a particular node type can control how 25 child nodes are visited during a preorder walk. (It can't control 26 the order during a postorder walk, because it is called _after_ 27 the walk has occurred.) The ASTVisitor modifies the visitor 28 argument by adding a visit method to the visitor; this method can 29 be used to visit a child node of arbitrary type. 30 """ 31 32 VERBOSE = 0 33 34 def __init__(self): 35 self.node = None 36 self._cache = {} 37 38 def default(self, node, *args): 39 for child in node.getChildNodes(): 40 self.dispatch(child, *args) 41 42 def dispatch(self, node, *args): 43 self.node = node 44 klass = node.__class__ 45 meth = self._cache.get(klass, None) 46 if meth is None: 47 className = klass.__name__ 48 meth = getattr(self.visitor, 'visit' + className, self.default) 49 self._cache[klass] = meth 50 ## if self.VERBOSE > 0: 51 ## className = klass.__name__ 52 ## if self.VERBOSE == 1: 53 ## if meth == 0: 54 ## print "dispatch", className 55 ## else: 56 ## print "dispatch", className, (meth and meth.__name__ or '') 57 return meth(node, *args) 58 59 def preorder(self, tree, visitor, *args): 60 """Do preorder walk of tree using visitor""" 61 self.visitor = visitor 62 visitor.visit = self.dispatch 63 self.dispatch(tree, *args) # XXX *args make sense? 64 65 class ExampleASTVisitor(ASTVisitor): 66 """Prints examples of the nodes that aren't visited 67 68 This visitor-driver is only useful for development, when it's 69 helpful to develop a visitor incrementally, and get feedback on what 70 you still have to do. 71 """ 72 examples = {} 73 74 def dispatch(self, node, *args): 75 self.node = node 76 meth = self._cache.get(node.__class__, None) 77 className = node.__class__.__name__ 78 if meth is None: 79 meth = getattr(self.visitor, 'visit' + className, 0) 80 self._cache[node.__class__] = meth 81 if self.VERBOSE > 1: 82 print "dispatch", className, (meth and meth.__name__ or '') 83 if meth: 84 meth(node, *args) 85 elif self.VERBOSE > 0: 86 klass = node.__class__ 87 if klass not in self.examples: 88 self.examples[klass] = klass 89 print 90 print self.visitor 91 print klass 92 for attr in dir(node): 93 if attr[0] != '_': 94 print "\t", "%-12.12s" % attr, getattr(node, attr) 95 print 96 return self.default(node, *args) 97 98 # XXX this is an API change 99 100 _walker = ASTVisitor 101 def walk(tree, visitor, walker=None, verbose=None): 102 if walker is None: 103 walker = _walker() 104 if verbose is not None: 105 walker.VERBOSE = verbose 106 walker.preorder(tree, visitor) 107 return walker.visitor 108 109 def dumpNode(node): 110 print node.__class__ 111 for attr in dir(node): 112 if attr[0] != '_': 113 print "\t", "%-10.10s" % attr, getattr(node, attr) 114
