diff --git a/lark/tree.py b/lark/tree.py index daf834a9..72bf6ac5 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -195,7 +195,6 @@ def expand_kids_by_data(self, *data_values): changed = True return changed - def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Leaf_T]: """Return all values in the tree that evaluate pred(value) as true. @@ -212,6 +211,20 @@ def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Le if pred(c): yield c + def replace_tokens(self, pred: 'Callable[[_Leaf_T], Optional[_Leaf_T]]') -> None: + """replace tokens in the tree using the result of pred(token) when it is not none. + + Example: + >>> tree.replace_tokens(lambda v: v.update(value=v.rstrip("\n")) if v.endswith("\n") else None) + """ + for index, child in enumerate(self.children): + if isinstance(child, Tree): + child.replace_tokens(pred) + else: + result = pred(child) + if result is not None: + self.children[index] = result + def __deepcopy__(self, memo): return type(self)(self.data, deepcopy(self.children, memo), meta=self._meta) diff --git a/tests/test_trees.py b/tests/test_trees.py index 82eceab1..b4f7244a 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -58,6 +58,16 @@ def test_find_token(self): tokens = list(self.tree2.find_token('T')) self.assertEqual(tokens, expected) + def test_replace_tokens(self): + tree2 = copy.deepcopy(self.tree2) + expected = Tree('a', [ + Tree('b', [Token('T', 'y')]), + Tree('c', [Token('T', 'y')]), + Tree('d', [Tree('z', [Token('T', 'zz'), Tree('zzz', 'zzz')])]), + ]) + tree2.replace_tokens(lambda v: v.update(value="y") if v == "x" else None) + self.assertEqual(tree2, expected) + def test_visitor(self): class Visitor1(Visitor): def __init__(self):