212 lines
7.8 KiB
Python
212 lines
7.8 KiB
Python
# -*- 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.
|
|
"""
|