def __init__(self) -> None:
CSTVisitor.__init__(self)
# List of gating matchers that we need to track and evaluate. We use these
# in conjuction with the call_if_inside and call_if_not_inside decorators
# to determine whether or not to call a visit/leave function.
self._matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]] = {
m: None for m in _gather_matchers(self)
}
# Mapping of matchers to functions. If in the course of visiting the tree,
# a node matches one of these matchers, the corresponding function will be
# called as if it was a visit_* method.
self._extra_visit_funcs: Dict[
BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
] = _gather_constructed_visit_funcs(self)
# Mapping of matchers to functions. If in the course of leaving the tree,
# a node matches one of these matchers, the corresponding function will be
# called as if it was a leave_* method.
self._extra_leave_funcs: Dict[
BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
] = _gather_constructed_leave_funcs(self)
# Make sure visit/leave functions constructed with @visit and @leave decorators
# have correct type annotations.
_check_types(
self._extra_visit_funcs,
"visit",
expected_param_count=1,
expected_none_return=True,
)
_check_types(
self._extra_leave_funcs,
"leave",
expected_param_count=1,
expected_none_return=True,
)
def on_leave_attribute(self, original_node: cst.CSTNode, attribute: str) -> None:
# Evaluate whether this current function has a decorator on it.
if _should_allow_visit(
self._matchers,
getattr(self, f"leave_{type(original_node).__name__}_{attribute}", None),
):
# Either the visit_func doesn't exist, we have no matchers, or we passed all
# matchers. In either case, just call the superclass behavior.
CSTVisitor.on_leave_attribute(self, original_node, attribute)
def on_leave(self, original_node: cst.CSTNode) -> None:
# First, evaluate whether this current function has a decorator on it.
if _should_allow_visit(
self._matchers, getattr(self, f"leave_{type(original_node).__name__}", None)
):
CSTVisitor.on_leave(self, original_node)
# Now, call any visitors that were hooked using a leave decorator.
for matcher, leave_funcs in reversed(list(self._extra_leave_funcs.items())):
if not self.matches(original_node, matcher):
continue
for leave_func in leave_funcs:
if _should_allow_visit(self._matchers, leave_func):
leave_func(original_node)
# Now, see if we have any matchers we should deactivate.
self._matchers = _leave_matchers(self._matchers, original_node)
def on_visit(self, node: cst.CSTNode) -> bool:
# First, evaluate any matchers that we have which we are not inside already.
self._matchers = _visit_matchers(self._matchers, node, self)
# Now, call any visitors that were hooked using a visit decorator.
_visit_constructed_funcs(self._extra_visit_funcs, self._matchers, node, self)
# Now, evaluate whether this current function has a decorator on it.
if not _should_allow_visit(
self._matchers, getattr(self, f"visit_{type(node).__name__}", None)
):
# We shouldn't visit this directly. However, we should continue
# visiting its children.
return True
# Either the visit_func doesn't exist, we have no matchers, or we passed all
# matchers. In either case, just call the superclass behavior.
return CSTVisitor.on_visit(self, node)
