From f222a7dee476f786244113ca8bb4037b69c2fba4 Mon Sep 17 00:00:00 2001 From: zzc <58017008+zzc0430@users.noreply.github.com> Date: Mon, 7 Apr 2025 16:09:30 +0800 Subject: [PATCH 1/9] add function `replace_nodes` to Tree This can be used to replace tokens in a `Tree` that satisfy a specified rule. --- lark/tree.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/lark/tree.py b/lark/tree.py index daf834a9..e201fdb4 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. @@ -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) From 6194e05b6175022de664c9d2d1c3e92f21a88713 Mon Sep 17 00:00:00 2001 From: zzc <58017008+zzc0430@users.noreply.github.com> Date: Tue, 8 Apr 2025 13:47:56 +0800 Subject: [PATCH 2/9] fix: Make sure the code logic is consistent with the function documentation. --- lark/tree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lark/tree.py b/lark/tree.py index e201fdb4..f627a502 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -222,7 +222,7 @@ def replace_nodes(self, pred: 'Callable[[Token], Optional[Token]]') -> None: self.replace_nodes(pred) elif isinstance(child, Token): result = pred(child) - if isinstance(result, Token): + if result is not None: self.children[index] = result def __deepcopy__(self, memo): From d529fbdcab194978bc7e0a7be9981cc180f2ef5a Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Tue, 8 Apr 2025 14:07:34 +0800 Subject: [PATCH 3/9] fix: fix traversal bug --- lark/tree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lark/tree.py b/lark/tree.py index f627a502..9d6415fc 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -219,7 +219,7 @@ def replace_nodes(self, pred: 'Callable[[Token], Optional[Token]]') -> None: """ for index, child in enumerate(self.children): if isinstance(child, Tree): - self.replace_nodes(pred) + child.replace_nodes(pred) elif isinstance(child, Token): result = pred(child) if result is not None: From 388fd80af56469302e2dac907cf13d973582df93 Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Tue, 8 Apr 2025 14:07:50 +0800 Subject: [PATCH 4/9] chore: add test --- tests/test_trees.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_trees.py b/tests/test_trees.py index 82eceab1..9406fb3a 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -58,6 +58,15 @@ def test_find_token(self): tokens = list(self.tree2.find_token('T')) self.assertEqual(tokens, expected) + def test_replace_nodes(self): + expected = Tree('a', [ + Tree('b', [Token('T', 'y')]), + Tree('c', [Token('T', 'y')]), + Tree('d', [Tree('z', [Token('T', 'zz'), Tree('zzz', 'zzz')])]), + ]) + self.tree2.replace_nodes(lambda v: v.update(value="y") if v == "x" else None) + self.assertEqual(self.tree2, expected) + def test_visitor(self): class Visitor1(Visitor): def __init__(self): From 545b8087ff521e536100a080f9b663c2175f6e7c Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Tue, 8 Apr 2025 14:21:24 +0800 Subject: [PATCH 5/9] refactor(tree): rename replace_nodes to replace_tokens for clarity --- lark/tree.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lark/tree.py b/lark/tree.py index 9d6415fc..4837de3a 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -211,15 +211,15 @@ def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Le if pred(c): yield c - def replace_nodes(self, pred: 'Callable[[Token], Optional[Token]]') -> None: + def replace_tokens(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) + >>> 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_nodes(pred) + child.replace_tokens(pred) elif isinstance(child, Token): result = pred(child) if result is not None: From caea7808abdc4dfff34d22b411ff736deb025279 Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Tue, 8 Apr 2025 14:22:34 +0800 Subject: [PATCH 6/9] fix: avoid changing `tree2` in-place --- tests/test_trees.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_trees.py b/tests/test_trees.py index 9406fb3a..36bee54c 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -58,14 +58,15 @@ def test_find_token(self): tokens = list(self.tree2.find_token('T')) self.assertEqual(tokens, expected) - def test_replace_nodes(self): + def test_replace_tokens(self): + tree2 = copy.copy(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')])]), ]) - self.tree2.replace_nodes(lambda v: v.update(value="y") if v == "x" else None) - self.assertEqual(self.tree2, expected) + 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): From 5ab7313081a8d7cc004796c7038591d7937ad17c Mon Sep 17 00:00:00 2001 From: zzc <58017008+zzc0430@users.noreply.github.com> Date: Tue, 8 Apr 2025 16:02:05 +0800 Subject: [PATCH 7/9] fix: use deepcopy in `test_replace_tokens` --- tests/test_trees.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_trees.py b/tests/test_trees.py index 36bee54c..b4f7244a 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -59,7 +59,7 @@ def test_find_token(self): self.assertEqual(tokens, expected) def test_replace_tokens(self): - tree2 = copy.copy(self.tree2) + tree2 = copy.deepcopy(self.tree2) expected = Tree('a', [ Tree('b', [Token('T', 'y')]), Tree('c', [Token('T', 'y')]), From c21139c95a074225f72486f85cf2b9e701830103 Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Wed, 9 Apr 2025 10:09:09 +0800 Subject: [PATCH 8/9] refactor(tree): update replace_tokens method to handle generic leaf types --- lark/tree.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lark/tree.py b/lark/tree.py index 4837de3a..1b94e4b3 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -211,7 +211,7 @@ def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Le if pred(c): yield c - def replace_tokens(self, pred: 'Callable[[Token], Optional[Token]]') -> None: + 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: @@ -220,7 +220,7 @@ def replace_tokens(self, pred: 'Callable[[Token], Optional[Token]]') -> None: for index, child in enumerate(self.children): if isinstance(child, Tree): child.replace_tokens(pred) - elif isinstance(child, Token): + else: result = pred(child) if result is not None: self.children[index] = result From b99ce77fd243fc0ce9015e60c27ed2c308a69683 Mon Sep 17 00:00:00 2001 From: zzc <1378113190@qq.com> Date: Wed, 9 Apr 2025 14:59:59 +0800 Subject: [PATCH 9/9] style: remove trailing whitespace in tree.py --- lark/tree.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lark/tree.py b/lark/tree.py index 1b94e4b3..72bf6ac5 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -210,7 +210,7 @@ def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Le else: 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. @@ -224,7 +224,7 @@ def replace_tokens(self, pred: 'Callable[[_Leaf_T], Optional[_Leaf_T]]') -> None 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)