KiBot/kibot/mcpy/visitors.py

59 lines
2.0 KiB
Python

# from functools import wraps
from ast import NodeTransformer, AST, copy_location, fix_missing_locations
from .unparse import unparse
class BaseMacroExpander(NodeTransformer):
"""
A base class for macro expander visitors. After identifying valid macro
syntax, the actual expander should return the result of calling `_expand()`
method with the proper arguments.
"""
def __init__(self, bindings):
self.bindings = bindings
def visit(self, tree):
"""Short-circuit visit() to avoid expansions if no macros."""
return super().visit(tree) if self.bindings else tree
def _expand(self, syntax, target, macroname, tree, kw=None):
"""
Transform `target` node, replacing it with the expansion result of
aplying the named macro on the proper node and recursively treat the
expansion as well.
"""
macro = self.bindings[macroname]
kw = kw or {}
kw.update({
'syntax': syntax,
'to_source': unparse,
'expand_macros': self.visit
})
expansion = _apply_macro(macro, tree, kw)
return self._visit_expansion(expansion, target)
def _visit_expansion(self, expansion, target):
"""
Ensures the macro expansions into None (deletions), other nodes or
list of nodes are expanded too.
"""
if expansion is not None:
is_node = isinstance(expansion, AST)
expansion = [expansion] if is_node else expansion
expansion = map(lambda n: copy_location(n, target), expansion)
expansion = map(fix_missing_locations, expansion)
expansion = map(self.visit, expansion)
expansion = list(expansion).pop() if is_node else list(expansion)
return expansion
def _ismacro(self, name):
return name in self.bindings
def _apply_macro(macro, tree, kw):
""" Executes the macro on tree passing extra kwargs. """
return macro(tree, **kw)