Home | History | Annotate | Download | only in compiler
      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