Lib/compiler/visitor.py

   1 from compiler import ast
   2
   3 class ASTVisitor:
   4     """Performs a depth-first walk of the AST
   5
   6     The ASTVisitor will walk the AST, performing either a preorder or
   7     postorder traversal depending on which method is called.
   8
   9     methods:
  10     preorder(tree, visitor)
  11     postorder(tree, visitor)
  12         tree: an instance of ast.Node
  13         visitor: an instance with visitXXX methods
  14
  15     The ASTVisitor is responsible for walking over the tree in the
  16     correct order.  For each node, it checks the visitor argument for
  17     a method named 'visitNodeType' where NodeType is the name of the
  18     node's class, e.g. Class.  If the method exists, it is called
  19     with the node as its sole argument.
  20
  21     The visitor method for a particular node type can control how
  22     child nodes are visited during a preorder walk.  (It can't control
  23     the order during a postorder walk, because it is called _after_
  24     the walk has occurred.)  The ASTVisitor modifies the visitor
  25     argument by adding a visit method to the visitor; this method can
  26     be used to visit a particular child node.  If the visitor method
  27     returns a true value, the ASTVisitor will not traverse the child
  28     nodes.
  29
  30     XXX The interface for controlling the preorder walk needs to be
  31     re-considered.  The current interface is convenient for visitors
  32     that mostly let the ASTVisitor do everything.  For something like
  33     a code generator, where you want to walk to occur in a specific
  34     order, it's a pain to add "return 1" to the end of each method.
  35     """
  36
  37     VERBOSE = 0
  38
  39     def __init__(self):
  40         self.node = None
  41         self._cache = {}
  42
  43     def preorder(self, tree, visitor):
  44         """Do preorder walk of tree using visitor"""
  45         self.visitor = visitor
  46         visitor.visit = self._preorder
  47         self._preorder(tree)
  48
  49     def _preorder(self, node, *args):
  50         return apply(self.dispatch, (node,) + args)
  51
  52     def default(self, node, *args):
  53         for child in node.getChildren():
  54             if isinstance(child, ast.Node):
  55                 apply(self._preorder, (child,) + args)
  56
  57     def dispatch(self, node, *args):
  58         self.node = node
  59         meth = self._cache.get(node.__class__, None)
  60         className = node.__class__.__name__
  61         if meth is None:
  62             meth = getattr(self.visitor, 'visit' + className, self.default)
  63             self._cache[node.__class__] = meth
  64         if self.VERBOSE > 0:
  65             if self.VERBOSE == 1:
  66                 if meth == 0:
  67                     print "dispatch", className
  68             else:
  69                 print "dispatch", className, (meth and meth.__name__ or '')
  70         return apply(meth, (node,) + args)
  71
  72 class ExampleASTVisitor(ASTVisitor):
  73     """Prints examples of the nodes that aren't visited
  74
  75     This visitor-driver is only useful for development, when it's
  76     helpful to develop a visitor incremently, and get feedback on what
  77     you still have to do.
  78     """
  79     examples = {}
  80
  81     def dispatch(self, node, *args):
  82         self.node = node
  83         meth = self._cache.get(node.__class__, None)
  84         className = node.__class__.__name__
  85         if meth is None:
  86             meth = getattr(self.visitor, 'visit' + className, 0)
  87             self._cache[node.__class__] = meth
  88         if self.VERBOSE > 1:
  89             print "dispatch", className, (meth and meth.__name__ or '')
  90         if meth:
  91             return apply(meth, (node,) + args)
  92         elif self.VERBOSE > 0:
  93             klass = node.__class__
  94             if not self.examples.has_key(klass):
  95                 self.examples[klass] = klass
  96                 print
  97                 print self.visitor
  98                 print klass
  99                 for attr in dir(node):
 100                     if attr[0] != '_':
 101                         print "\t", "%-12.12s" % attr, getattr(node, attr)
 102                 print
 103             return apply(self.default, (node,) + args)
 104
 105 _walker = ASTVisitor
 106 def walk(tree, visitor, verbose=None):
 107     w = _walker()
 108     if verbose is not None:
 109         w.VERBOSE = verbose
 110     w.preorder(tree, visitor)
 111     return w.visitor
 112
 113 def dumpNode(node):
 114     print node.__class__
 115     for attr in dir(node):
 116         if attr[0] != '_':
 117             print "\t", "%-10.10s" % attr, getattr(node, attr)