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