# -*- coding: utf-8; -*- """AST walkers. These have a state stack and a node collector, and accept also a `list` of statements. Otherwise they work like `ast.NodeVisitor` and `ast.NodeTransformer`. Basic usage summary:: def kittify(mytree): class Kittifier(ASTTransformer): def transform(self, tree): if type(tree) is ast.Constant: self.collect(tree.value) tree.value = "meow!" if self.state.meows % 2 == 0 else "miaow!" self.state.meows += 1 return self.generic_visit(tree) # recurse w = Kittifier(meows=0) # set the initial state here mytree = w.visit(mytree) # it's basically an ast.NodeTransformer print(w.collected) # collected values, in the order visited return mytree def getmeows(mytree): class MeowCollector(ASTVisitor): def examine(self, tree): if type(tree) is ast.Constant and tree.value in ("meow!", "miaow!"): self.collect(tree) self.generic_visit(tree) w = MeowCollector() w.visit(mytree) print(w.collected) return w.collected For an example of `withstate`, see `mcpyrate.astfixers`. """ __all__ = ["ASTVisitor", "ASTTransformer"] from abc import ABCMeta, abstractmethod from ast import NodeVisitor, NodeTransformer, iter_child_nodes from .bunch import Bunch from . import utils class BaseASTWalker: """AST walker base class, providing a state stack and a node collector.""" def __init__(self, **bindings): """Bindings are loaded into the initial `self.state` as attributes.""" self.reset(**bindings) def reset(self, **bindings): """Clear everything. Load new bindings into a blank `self.state`.""" self._stack = [Bunch(**bindings)] self._subtree_overrides = {} self.collected = [] def _setstate(self, newstate): self._stack[-1] = newstate return newstate def _getstate(self): return self._stack[-1] state = property(fget=_getstate, fset=_setstate, doc="The current state. Mutable. Can be rebound to replace it.") def withstate(self, tree, **bindings): """Arrange to visit a subtree with a temporarily replaced, updated state. `tree` can be an AST node or a statement suite (`list` of AST nodes). It is identified by `id(tree)` at enter time. Bindings update a copy of `self.state`. If several `withstate` calls are made for the same `tree`, the last one overrides. Generally speaking: `withstate(subtree, ...)` should be used if you then intend to `visit(subtree)`, which recurses into that node (or suite) only. `generic_withstate(tree, ...)` should be used if you then intend to `generic_visit(tree)`, which recurses into the children of `tree`. (It is possible to mix and match, but think through what you're doing.) """ newstate = self.state.copy() newstate.update(**bindings) # Due to how `ast.NodeTransformer.generic_visit` works, `visit` is # never called for a statement suite, but only separately for each # individual statement in it. So if we should set `newstate` for a # suite, do it for each statement in it. if isinstance(tree, list): for elt in tree: self._subtree_overrides[id(elt)] = newstate else: self._subtree_overrides[id(tree)] = newstate def generic_withstate(self, tree, **bindings): """Like `withstate`, but set up the new state for all children of `tree`. The same state instance is shared between the child nodes. If several `generic_withstate` calls are made for the same `tree`, the last one overrides (assuming the list of children has not changed in between). The silly name is because this relates to `withstate` as `generic_visit` relates to `visit`. Generally speaking: `generic_withstate(tree, ...)` should be used if you then intend to `generic_visit(tree)`, which recurses into the children of `tree`. `withstate(subtree, ...)` should be used if you then intend to `visit(subtree)`, which recurses into that node (or suite) only. (It is possible to mix and match, but think through what you're doing.) """ newstate = self.state.copy() newstate.update(**bindings) for node in iter_child_nodes(tree): self._subtree_overrides[id(node)] = newstate def collect(self, value): """Collect a value. The values are placed in the list `self.collected`.""" self.collected.append(value) return value class ASTVisitor(BaseASTWalker, NodeVisitor, metaclass=ABCMeta): """AST visitor, like `ast.NodeVisitor`, but with the features of `BaseASTWalker`. (Because `return tree`, or `return self.generic_visit(tree)`, is too much to remember when not editing.) If you want to edit the tree, use `ASTTransformer` instead. """ def visit(self, tree): """Start visiting `tree`. **Do not override this method; see `examine` instead.**""" newstate = self._subtree_overrides.pop(id(tree), False) if newstate: self._stack.append(newstate) try: if isinstance(tree, list): for elt in tree: self.visit(elt) return return self.examine(tree) finally: if newstate: self._stack.pop() @abstractmethod def examine(self, tree): """Examine one node. **Abstract method, override this.** There is only one `examine` method. To detect node type, use `type(tree)`. This method must recurse explicitly where needed. Use: - `self.generic_visit(tree)` to visit all children of `tree`. - `self.visit(tree.something)` to selectively visit only some children. Visiting a statement suite with `self.visit` is also ok; it is treated like a `generic_visit`. As in `ast.NodeVisitor`: - Return value of `examine` is forwarded by `visit`. - `generic_visit` always returns `None`. """ class ASTTransformer(BaseASTWalker, NodeTransformer, metaclass=ABCMeta): """AST transformer, like `ast.NodeTransformer`, but with the features of `BaseASTWalker`. If you only want to examine the tree, not edit it, consider `ASTVisitor` instead. """ def visit(self, tree): """Start transforming `tree`. **Do not override this method; see `transform` instead.**""" newstate = self._subtree_overrides.pop(id(tree), False) if newstate: self._stack.append(newstate) try: if isinstance(tree, list): new_tree = utils.flatten(self.visit(elt) for elt in tree) if not new_tree: new_tree = [] # preserve the type of `tree`; an empty list shouldn't turn into `None` tree[:] = new_tree return tree return self.transform(tree) finally: if newstate: self._stack.pop() @abstractmethod def transform(self, tree): """Transform one node. **Abstract method, override this.** There is only one `transform` method. To detect node type, use `type(tree)`. This method must recurse explicitly where needed. Use: - `tree = self.generic_visit(tree)` to visit all children of `tree`. - `tree.something = self.visit(tree.something)` to selectively visit only some children. Visiting a statement suite with `self.visit` is also ok. Return value as in `ast.NodeTransformer`. If you don't want to make changes, you must `return tree`. If you return `None`, the subtree is removed. """