175 lines
5.1 KiB
Python
175 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from collections.abc import Mapping
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
|
|
from .objectmodel import Node
|
|
from .util import is_list
|
|
|
|
|
|
class NodeWalkerMeta(type):
|
|
def __new__(mcs, name, bases, dct):
|
|
class_ = super().__new__(mcs, name, bases, dct)
|
|
class_._walker_cache = {}
|
|
return class_
|
|
|
|
|
|
class NodeWalker(metaclass=NodeWalkerMeta):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# copy the class attribute to avoid linter warnings
|
|
self._walker_cache = type(
|
|
self,
|
|
)._walker_cache # pylint: disable=no-member
|
|
|
|
def walk(self, node: Node | list[Node], *args, **kwargs) -> Any:
|
|
if isinstance(node, list | tuple):
|
|
return [self.walk(n, *args, **kwargs) for n in node]
|
|
|
|
if isinstance(node, Mapping):
|
|
return {
|
|
name: self.walk(value, *args, **kwargs)
|
|
for name, value in node.items()
|
|
}
|
|
|
|
walker = self._find_walker(node)
|
|
if callable(walker):
|
|
return walker(self, node, *args, **kwargs)
|
|
else:
|
|
return node
|
|
|
|
def walk_children(self, node: Node, *args, **kwargs):
|
|
if not isinstance(node, Node):
|
|
return []
|
|
|
|
return [
|
|
self.walk(child, *args, **kwargs)
|
|
for child in node.children()
|
|
]
|
|
|
|
# note: backwards compatibility
|
|
_walk_children = walk_children
|
|
|
|
def _find_walker(self, node: Node, prefix='walk_'):
|
|
def pythonize_match(m):
|
|
return '_' + m.group().lower()
|
|
|
|
cls = self.__class__
|
|
node_cls = node.__class__
|
|
node_cls_qualname = node_cls.__qualname__
|
|
|
|
if walker := self._walker_cache.get(node_cls_qualname):
|
|
return walker
|
|
|
|
node_classes = [node.__class__]
|
|
while node_classes:
|
|
node_cls = node_classes.pop(0)
|
|
|
|
cammelcase_name = node_cls.__name__
|
|
walker = getattr(cls, prefix + cammelcase_name, None)
|
|
if callable(walker):
|
|
break
|
|
|
|
# walk__pythonic_name with double underscore after walk
|
|
pythonic_name = re.sub(
|
|
r'[A-Z]+', pythonize_match, node_cls.__name__,
|
|
)
|
|
if pythonic_name != cammelcase_name:
|
|
walker = getattr(cls, prefix + pythonic_name, None)
|
|
if callable(walker):
|
|
break
|
|
|
|
# walk_pythonic_name with single underscore after prefix
|
|
pythonic_name = pythonic_name.lstrip('_')
|
|
if pythonic_name != cammelcase_name:
|
|
walker = getattr(cls, prefix + pythonic_name, None)
|
|
if callable(walker):
|
|
break
|
|
|
|
for b in node_cls.__bases__:
|
|
if b not in node_classes:
|
|
node_classes.append(b)
|
|
else:
|
|
walker = getattr(cls, '_walk_default', None)
|
|
if walker is None:
|
|
walker = getattr(
|
|
cls, 'walk_default', None,
|
|
) # backwards compatibility
|
|
if not callable(walker):
|
|
walker = None
|
|
|
|
self._walker_cache[node_cls_qualname] = walker
|
|
return walker
|
|
|
|
|
|
class PreOrderWalker(NodeWalker):
|
|
def walk(self, node, *args, **kwargs):
|
|
result = super().walk(node, *args, **kwargs)
|
|
if result is not None:
|
|
self.walk_children(node, *args, **kwargs)
|
|
return result
|
|
|
|
|
|
class DepthFirstWalker(NodeWalker):
|
|
def walk(self, node, *args, **kwargs):
|
|
if isinstance(node, Node):
|
|
children = [self.walk(c, *args, **kwargs) for c in node.children()]
|
|
return super().walk(node, children, *args, **kwargs)
|
|
elif isinstance(node, Mapping):
|
|
return {n: self.walk(e, *args, **kwargs) for n, e in node.items()}
|
|
elif is_list(node):
|
|
return [self.walk(e, *args, **kwargs) for e in iter(node)]
|
|
else:
|
|
return super().walk(node, [], *args, **kwargs)
|
|
|
|
|
|
class ContextWalker(NodeWalker):
|
|
def __init__(self, initial_context):
|
|
super().__init__()
|
|
self._initial_context = initial_context
|
|
self._context_stack = [initial_context]
|
|
|
|
# abstract
|
|
def get_node_context(self, node, *args, **kwargs):
|
|
return node
|
|
|
|
# abstract
|
|
def enter_context(self, ctx):
|
|
pass
|
|
|
|
# abstract
|
|
def leave_context(self, ctx):
|
|
pass
|
|
|
|
def push_context(self, ctx):
|
|
self._context_stack.append(ctx)
|
|
|
|
def pop_context(self):
|
|
self._context_stack.pop()
|
|
|
|
@property
|
|
def initial_context(self):
|
|
return self._initial_context
|
|
|
|
@property
|
|
def context(self):
|
|
return self._context_stack[-1]
|
|
|
|
@contextmanager
|
|
def new_context(self, node, *args, **kwargs):
|
|
ctx = self.get_node_context(node, *args, **kwargs)
|
|
if ctx == self.context:
|
|
yield ctx
|
|
else:
|
|
self.enter_context(ctx)
|
|
try:
|
|
self.push_context(ctx)
|
|
try:
|
|
yield ctx
|
|
finally:
|
|
self.pop_context()
|
|
finally:
|
|
self.leave_context(ctx)
|