1from compiler import ast
2
3# XXX should probably rename ASTVisitor to ASTWalker
4# XXX can it be made even more generic?
5
6class 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
65class 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
101def 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
109def 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