|
| 1 | +from tools 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 | + stop = apply(self.dispatch, (node,) + args) |
| 51 | + if stop: |
| 52 | + return |
| 53 | + for child in node.getChildren(): |
| 54 | + if isinstance(child, ast.Node): |
| 55 | + self._preorder(child) |
| 56 | + |
| 57 | + def postorder(self, tree, visitor): |
| 58 | + """Do preorder walk of tree using visitor""" |
| 59 | + self.visitor = visitor |
| 60 | + visitor.visit = self._postorder |
| 61 | + self._postorder(tree) |
| 62 | + |
| 63 | + def _postorder(self, tree, *args): |
| 64 | + for child in node.getChildren(): |
| 65 | + if isinstance(child, ast.Node): |
| 66 | + self._postorder(child) |
| 67 | + apply(self.dispatch, (node,) + args) |
| 68 | + |
| 69 | + def dispatch(self, node, *args): |
| 70 | + self.node = node |
| 71 | + meth = self._cache.get(node.__class__, None) |
| 72 | + className = node.__class__.__name__ |
| 73 | + if meth is None: |
| 74 | + meth = getattr(self.visitor, 'visit' + className, 0) |
| 75 | + self._cache[node.__class__] = meth |
| 76 | + if self.VERBOSE > 0: |
| 77 | + if self.VERBOSE == 1: |
| 78 | + if meth == 0: |
| 79 | + print "dispatch", className |
| 80 | + else: |
| 81 | + print "dispatch", className, (meth and meth.__name__ or '') |
| 82 | + if meth: |
| 83 | + return apply(meth, (node,) + args) |
| 84 | + |
| 85 | +class ExampleASTVisitor(ASTVisitor): |
| 86 | + """Prints examples of the nodes that aren't visited |
| 87 | +
|
| 88 | + This visitor-driver is only useful for development, when it's |
| 89 | + helpful to develop a visitor incremently, and get feedback on what |
| 90 | + you still have to do. |
| 91 | + """ |
| 92 | + examples = {} |
| 93 | + |
| 94 | + def dispatch(self, node, *args): |
| 95 | + self.node = node |
| 96 | + className = node.__class__.__name__ |
| 97 | + meth = getattr(self.visitor, 'visit' + className, None) |
| 98 | + if self.VERBOSE > 0: |
| 99 | + if self.VERBOSE > 1: |
| 100 | + print "dispatch", className, (meth and meth.__name__ or '') |
| 101 | + if meth: |
| 102 | + return apply(meth, (node,) + args) |
| 103 | + elif self.VERBOSE > 0: |
| 104 | + klass = node.__class__ |
| 105 | + if self.examples.has_key(klass): |
| 106 | + return |
| 107 | + self.examples[klass] = klass |
| 108 | + print |
| 109 | + print klass |
| 110 | + for attr in dir(node): |
| 111 | + if attr[0] != '_': |
| 112 | + print "\t", "%-12.12s" % attr, getattr(node, attr) |
| 113 | + print |
| 114 | + |
| 115 | +_walker = ASTVisitor |
| 116 | +def walk(tree, visitor, verbose=None): |
| 117 | + w = _walker() |
| 118 | + if verbose is not None: |
| 119 | + w.VERBOSE = verbose |
| 120 | + w.preorder(tree, visitor) |
| 121 | + return w.visitor |
| 122 | + |
| 123 | +def dumpNode(node): |
| 124 | + print node.__class__ |
| 125 | + for attr in dir(node): |
| 126 | + if attr[0] != '_': |
| 127 | + print "\t", "%-10.10s" % attr, getattr(node, attr) |
0 commit comments