add function replace_nodes to Tree

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)
 

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):