diff --git a/lark/tree.py b/lark/tree.py index d496d75..e6d5ed7 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -174,6 +174,30 @@ class Visitor_NoRecurse(Visitor): return tree +from functools import wraps +def visit_children_decor(func): + @wraps(func) + def inner(cls, tree): + values = cls.visit_children(tree) + return func(cls, values) + return inner + +class Interpreter(object): + + def visit(self, tree): + return getattr(self, tree.data)(tree) + + def visit_children(self, tree): + return [self.visit(child) if isinstance(child, Tree) else child + for child in tree.children] + + def __getattr__(self, name): + return self.__default__ + + def __default__(self, tree): + self.visit_children(tree) + + class Transformer_NoRecurse(Transformer): def transform(self, tree): subtrees = list(tree.iter_subtrees()) diff --git a/tests/test_trees.py b/tests/test_trees.py index c90cc7d..c83b5ef 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -5,7 +5,7 @@ from unittest import TestCase import copy import pickle -from lark.tree import Tree +from lark.tree import Tree, Interpreter, visit_children_decor class TestTrees(TestCase): @@ -21,6 +21,37 @@ class TestTrees(TestCase): assert pickle.loads(data) == s + def test_interp(self): + t = Tree('a', [Tree('b', []), Tree('c', []), 'd']) + + class Interp1(Interpreter): + def a(self, tree): + return self.visit_children(tree) + ['e'] + + def b(self, tree): + return 'B' + + def c(self, tree): + return 'C' + + self.assertEqual(Interp1().visit(t), list('BCde')) + + class Interp2(Interpreter): + @visit_children_decor + def a(self, values): + return values + ['e'] + + def b(self, tree): + return 'B' + + def c(self, tree): + return 'C' + + self.assertEqual(Interp2().visit(t), list('BCde')) + + + + if __name__ == '__main__': unittest.main()