Lib/compiler/visitor.py

   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 particular child node.  If the visitor method
  30     returns a true value, the ASTVisitor will not traverse the child
  31     nodes.
  32
  33     XXX The interface for controlling the preorder walk needs to be
  34     re-considered.  The current interface is convenient for visitors
  35     that mostly let the ASTVisitor do everything.  For something like
  36     a code generator, where you want to walk to occur in a specific
  37     order, it's a pain to add "return 1" to the end of each method.
  38     """
  39
  40     VERBOSE = 0
  41
  42     def __init__(self):
  43         self.node = None
  44         self._cache = {}
  45
  46     def default(self, node, *args):
  47         for child in node.getChildNodes():
  48             self.dispatch(child, *args)
  49
  50     def dispatch(self, node, *args):
  51         self.node = node
  52         klass = node.__class__
  53         meth = self._cache.get(klass, None)
  54         if meth is None:
  55             className = klass.__name__
  56             meth = getattr(self.visitor, 'visit' + className, self.default)
  57             self._cache[klass] = meth
  58 ##        if self.VERBOSE > 0:
  59 ##            className = klass.__name__
  60 ##            if self.VERBOSE == 1:
  61 ##                if meth == 0:
  62 ##                    print "dispatch", className
  63 ##            else:
  64 ##                print "dispatch", className, (meth and meth.__name__ or '')
  65         return meth(node, *args)
  66
  67     def preorder(self, tree, visitor, *args):
  68         """Do preorder walk of tree using visitor"""
  69         self.visitor = visitor
  70         visitor.visit = self.dispatch
  71         self.dispatch(tree, *args) # XXX *args make sense?
  72
  73 class ExampleASTVisitor(ASTVisitor):
  74     """Prints examples of the nodes that aren't visited
  75
  76     This visitor-driver is only useful for development, when it's
  77     helpful to develop a visitor incremently, and get feedback on what
  78     you still have to do.
  79     """
  80     examples = {}
  81
  82     def dispatch(self, node, *args):
  83         self.node = node
  84         meth = self._cache.get(node.__class__, None)
  85         className = node.__class__.__name__
  86         if meth is None:
  87             meth = getattr(self.visitor, 'visit' + className, 0)
  88             self._cache[node.__class__] = meth
  89         if self.VERBOSE > 1:
  90             print "dispatch", className, (meth and meth.__name__ or '')
  91         if meth:
  92             meth(node, *args)
  93         elif self.VERBOSE > 0:
  94             klass = node.__class__
  95             if not self.examples.has_key(klass):
  96                 self.examples[klass] = klass
  97                 print
  98                 print self.visitor
  99                 print klass
 100                 for attr in dir(node):
 101                     if attr[0] != '_':
 102                         print "\t", "%-12.12s" % attr, getattr(node, attr)
 103                 print
 104             return self.default(node, *args)
 105
 106 # XXX this is an API change
 107
 108 _walker = ASTVisitor
 109 def walk(tree, visitor, walker=None, verbose=None):
 110     if walker is None:
 111         walker = _walker()
 112     if verbose is not None:
 113         walker.VERBOSE = verbose
 114     walker.preorder(tree, visitor)
 115     return walker.visitor
 116
 117 def dumpNode(node):
 118     print node.__class__
 119     for attr in dir(node):
 120         if attr[0] != '_':
 121             print "\t", "%-10.10s" % attr, getattr(node, attr)