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.
 
@@ -211,7 +210,21 @@ def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Le
             else:
                 if pred(c):
                     yield c
+
+    def replace_nodes(self, pred: 'Callable[[Token], Optional[Token]]') -> None:
+        """replace tokens in the tree using the result of pred(token) when it is not none.
 
+        Example:
+            >>> tree.replace_nodes(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):
+                self.replace_nodes(pred)
+            elif isinstance(child, Token):
+                result = pred(child)
+                if isinstance(result, Token):
+                    self.children[index] = result
+
     def __deepcopy__(self, memo):
         return type(self)(self.data, deepcopy(self.children, memo), meta=self._meta)