factor out the tree walking/visitor code that was in compile.py

1 files changed, 127 insertions, 0 deletions

diff --git a/Lib/compiler/visitor.py b/Lib/compiler/visitor.py
new file mode 100644
index 0000000..5bdf108
--- /dev/null
+++ b/Lib/compiler/visitor.py
@@ -0,0 +1,127 @@
+from tools import ast
+
+class ASTVisitor:
+    """Performs a depth-first walk of the AST
+
+    The ASTVisitor will walk the AST, performing either a preorder or
+    postorder traversal depending on which method is called.
+
+    methods:
+    preorder(tree, visitor)
+    postorder(tree, visitor)
+        tree: an instance of ast.Node
+        visitor: an instance with visitXXX methods
+
+    The ASTVisitor is responsible for walking over the tree in the
+    correct order.  For each node, it checks the visitor argument for
+    a method named 'visitNodeType' where NodeType is the name of the
+    node's class, e.g. Class.  If the method exists, it is called
+    with the node as its sole argument.
+
+    The visitor method for a particular node type can control how
+    child nodes are visited during a preorder walk.  (It can't control
+    the order during a postorder walk, because it is called _after_
+    the walk has occurred.)  The ASTVisitor modifies the visitor
+    argument by adding a visit method to the visitor; this method can
+    be used to visit a particular child node.  If the visitor method
+    returns a true value, the ASTVisitor will not traverse the child
+    nodes.
+
+    XXX The interface for controlling the preorder walk needs to be
+    re-considered.  The current interface is convenient for visitors
+    that mostly let the ASTVisitor do everything.  For something like
+    a code generator, where you want to walk to occur in a specific
+    order, it's a pain to add "return 1" to the end of each method.
+    """
+
+    VERBOSE = 0
+
+    def __init__(self):
+        self.node = None
+	self._cache = {}
+
+    def preorder(self, tree, visitor):
+        """Do preorder walk of tree using visitor"""
+        self.visitor = visitor
+        visitor.visit = self._preorder
+        self._preorder(tree)
+
+    def _preorder(self, node, *args):
+	stop = apply(self.dispatch, (node,) + args)
+        if stop:
+            return
+        for child in node.getChildren():
+            if isinstance(child, ast.Node):
+                self._preorder(child)
+
+    def postorder(self, tree, visitor):
+        """Do preorder walk of tree using visitor"""
+        self.visitor = visitor
+        visitor.visit = self._postorder
+        self._postorder(tree)
+
+    def _postorder(self, tree, *args):
+        for child in node.getChildren():
+            if isinstance(child, ast.Node):
+                self._postorder(child)
+	apply(self.dispatch, (node,) + args)
+
+    def dispatch(self, node, *args):
+        self.node = node
+	meth = self._cache.get(node.__class__, None)
+	className = node.__class__.__name__
+	if meth is None:
+	    meth = getattr(self.visitor, 'visit' + className, 0)
+	    self._cache[node.__class__] = meth
+        if self.VERBOSE > 0:
+            if self.VERBOSE == 1:
+                if meth == 0:
+                    print "dispatch", className
+            else:
+                print "dispatch", className, (meth and meth.__name__ or '')
+        if meth:
+            return apply(meth, (node,) + args)
+
+class ExampleASTVisitor(ASTVisitor):
+    """Prints examples of the nodes that aren't visited
+
+    This visitor-driver is only useful for development, when it's
+    helpful to develop a visitor incremently, and get feedback on what
+    you still have to do.
+    """
+    examples = {}
+    
+    def dispatch(self, node, *args):
+        self.node = node
+        className = node.__class__.__name__
+        meth = getattr(self.visitor, 'visit' + className, None)
+        if self.VERBOSE > 0:
+            if self.VERBOSE > 1:
+                print "dispatch", className, (meth and meth.__name__ or '')
+        if meth:
+            return apply(meth, (node,) + args)
+        elif self.VERBOSE > 0:
+            klass = node.__class__
+            if self.examples.has_key(klass):
+                return
+            self.examples[klass] = klass
+            print
+            print klass
+            for attr in dir(node):
+                if attr[0] != '_':
+                    print "\t", "%-12.12s" % attr, getattr(node, attr)
+            print
+
+_walker = ASTVisitor
+def walk(tree, visitor, verbose=None):
+    w = _walker()
+    if verbose is not None:
+        w.VERBOSE = verbose
+    w.preorder(tree, visitor)
+    return w.visitor
+
+def dumpNode(node):
+    print node.__class__
+    for attr in dir(node):
+        if attr[0] != '_':
+            print "\t", "%-10.10s" % attr, getattr(node, attr)