| 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 not self.examples.has_key(klass):
|
|---|
| 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)
|
|---|
