10a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom compiler import ast 20a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 30a8c90248264a8b26970b4473770bcc3df8515fJosh Gao# XXX should probably rename ASTVisitor to ASTWalker 40a8c90248264a8b26970b4473770bcc3df8515fJosh Gao# XXX can it be made even more generic? 50a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 60a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass ASTVisitor: 70a8c90248264a8b26970b4473770bcc3df8515fJosh Gao """Performs a depth-first walk of the AST 80a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 90a8c90248264a8b26970b4473770bcc3df8515fJosh Gao The ASTVisitor will walk the AST, performing either a preorder or 100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao postorder traversal depending on which method is called. 110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao methods: 130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao preorder(tree, visitor) 140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao postorder(tree, visitor) 150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao tree: an instance of ast.Node 160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao visitor: an instance with visitXXX methods 170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao The ASTVisitor is responsible for walking over the tree in the 190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao correct order. For each node, it checks the visitor argument for 200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao a method named 'visitNodeType' where NodeType is the name of the 210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao node's class, e.g. Class. If the method exists, it is called 220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao with the node as its sole argument. 230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao The visitor method for a particular node type can control how 250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao child nodes are visited during a preorder walk. (It can't control 260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao the order during a postorder walk, because it is called _after_ 270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao the walk has occurred.) The ASTVisitor modifies the visitor 280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao argument by adding a visit method to the visitor; this method can 290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao be used to visit a child node of arbitrary type. 300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao """ 310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao VERBOSE = 0 330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao def __init__(self): 350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.node = None 360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self._cache = {} 370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao def default(self, node, *args): 390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao for child in node.getChildNodes(): 400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.dispatch(child, *args) 410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao def dispatch(self, node, *args): 430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.node = node 440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao klass = node.__class__ 450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao meth = self._cache.get(klass, None) 460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if meth is None: 470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao className = klass.__name__ 480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao meth = getattr(self.visitor, 'visit' + className, self.default) 490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self._cache[klass] = meth 500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## if self.VERBOSE > 0: 510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## className = klass.__name__ 520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## if self.VERBOSE == 1: 530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## if meth == 0: 540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## print "dispatch", className 550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## else: 560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao## print "dispatch", className, (meth and meth.__name__ or '') 570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao return meth(node, *args) 580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao def preorder(self, tree, visitor, *args): 600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao """Do preorder walk of tree using visitor""" 610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.visitor = visitor 620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao visitor.visit = self.dispatch 630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.dispatch(tree, *args) # XXX *args make sense? 640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 650a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass ExampleASTVisitor(ASTVisitor): 660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao """Prints examples of the nodes that aren't visited 670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao This visitor-driver is only useful for development, when it's 690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao helpful to develop a visitor incrementally, and get feedback on what 700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao you still have to do. 710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao """ 720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao examples = {} 730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao def dispatch(self, node, *args): 750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.node = node 760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao meth = self._cache.get(node.__class__, None) 770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao className = node.__class__.__name__ 780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if meth is None: 790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao meth = getattr(self.visitor, 'visit' + className, 0) 800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self._cache[node.__class__] = meth 810a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if self.VERBOSE > 1: 820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print "dispatch", className, (meth and meth.__name__ or '') 830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if meth: 840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao meth(node, *args) 850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao elif self.VERBOSE > 0: 860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao klass = node.__class__ 870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if klass not in self.examples: 880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao self.examples[klass] = klass 890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print 900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print self.visitor 910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print klass 920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao for attr in dir(node): 930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if attr[0] != '_': 940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print "\t", "%-12.12s" % attr, getattr(node, attr) 950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print 960a8c90248264a8b26970b4473770bcc3df8515fJosh Gao return self.default(node, *args) 970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao# XXX this is an API change 990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 1000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao_walker = ASTVisitor 1010a8c90248264a8b26970b4473770bcc3df8515fJosh Gaodef walk(tree, visitor, walker=None, verbose=None): 1020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if walker is None: 1030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao walker = _walker() 1040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if verbose is not None: 1050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao walker.VERBOSE = verbose 1060a8c90248264a8b26970b4473770bcc3df8515fJosh Gao walker.preorder(tree, visitor) 1070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao return walker.visitor 1080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao 1090a8c90248264a8b26970b4473770bcc3df8515fJosh Gaodef dumpNode(node): 1100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print node.__class__ 1110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao for attr in dir(node): 1120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao if attr[0] != '_': 1130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao print "\t", "%-10.10s" % attr, getattr(node, attr) 114