Lib/compiler/visitor.py

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