From bd40098f1d3ef10b7201aa9b81f9960e2ecdb9be Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Wed, 11 Sep 2024 23:08:01 +0300 Subject: [PATCH 1/6] Add Tree.find_token() method --- lark/tree.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/lark/tree.py b/lark/tree.py index 76f8738e..41796c7d 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -171,6 +171,10 @@ def find_data(self, data: str) -> 'Iterator[Tree[_Leaf_T]]': ###} + def find_token(self, typ: str) -> Iterator[_Leaf_T]: + """Returns all tokens whose type equals the given typ.""" + return self.scan_values(lambda v: isinstance(v, Token) and v.type == typ) + def expand_kids_by_data(self, *data_values): """Expand (inline) children with any of the given data values. Returns True if anything changed""" changed = False From 2472667a73d28c2c4bb16a5610661636e69d5046 Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Thu, 12 Sep 2024 08:08:22 +0300 Subject: [PATCH 2/6] Add test --- tests/test_trees.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_trees.py b/tests/test_trees.py index 1f69869e..49a53b32 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -17,6 +17,11 @@ class TestTrees(TestCase): def setUp(self): self.tree1 = Tree('a', [Tree(x, y) for x, y in zip('bcd', 'xyz')]) + self.tree2 = Tree('a', [ + Tree('b', [Token('T', 'x')]), + Tree('c', [Token('T', 'y')]), + Tree('d', [Tree('z', [Token('T', 'zz')])]), + ]) def test_eq(self): assert self.tree1 == self.tree1 @@ -48,6 +53,11 @@ def test_iter_subtrees_topdown(self): nodes = list(self.tree1.iter_subtrees_topdown()) self.assertEqual(nodes, expected) + def test_find_token(self): + expected = [Token('T', 'x'), Token('T', 'y'), Token('T', 'zz')] + tokens = list(self.tree2.find_token('T')) + self.assertEqual(tokens, expected) + def test_visitor(self): class Visitor1(Visitor): def __init__(self): From 979f6e29f68a6d20b8b97eb9e8916ccb0bef1333 Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Thu, 12 Sep 2024 08:09:21 +0300 Subject: [PATCH 3/6] Update .gitignore to exclude artefacts generated by tox --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 26275e4b..d4e64180 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ *.pyc *.pyo /.tox +/lark.egg-info/** /lark_parser.egg-info/** tags .vscode From 01c4f42ba556400e89c0b4ce105d2138cd2d805d Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Thu, 12 Sep 2024 08:10:40 +0300 Subject: [PATCH 4/6] Fix bug, add docs --- lark/tree.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lark/tree.py b/lark/tree.py index 41796c7d..321c6455 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -171,9 +171,15 @@ def find_data(self, data: str) -> 'Iterator[Tree[_Leaf_T]]': ###} - def find_token(self, typ: str) -> Iterator[_Leaf_T]: - """Returns all tokens whose type equals the given typ.""" - return self.scan_values(lambda v: isinstance(v, Token) and v.type == typ) + def find_token(self, token_type: str) -> Iterator[_Leaf_T]: + """Returns all tokens whose type equals the given token_type. + + This is a recursive function that will find tokens in all the subtrees. + + Example: + >>> term_tokens = tree.find_token('TERM') + """ + return self.scan_values(lambda v: not isinstance(v, Tree) and v.type == token_type) def expand_kids_by_data(self, *data_values): """Expand (inline) children with any of the given data values. Returns True if anything changed""" From c48b7d38c14e9c581f6e169bfa898c35c4e416bd Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Fri, 3 Jan 2025 08:15:19 +0200 Subject: [PATCH 5/6] Use `isinstance(v, Token)` --- lark/tree.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lark/tree.py b/lark/tree.py index 321c6455..9dccadd7 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -3,8 +3,10 @@ from typing import List, Callable, Iterator, Union, Optional, Generic, TypeVar, TYPE_CHECKING +from .lexer import Token + if TYPE_CHECKING: - from .lexer import TerminalDef, Token + from .lexer import TerminalDef try: import rich except ImportError: @@ -179,7 +181,7 @@ def find_token(self, token_type: str) -> Iterator[_Leaf_T]: Example: >>> term_tokens = tree.find_token('TERM') """ - return self.scan_values(lambda v: not isinstance(v, Tree) and v.type == token_type) + return self.scan_values(lambda v: isinstance(v, Token) and v.type == token_type) def expand_kids_by_data(self, *data_values): """Expand (inline) children with any of the given data values. Returns True if anything changed""" From 971e418bd5ff4e0889ad8b6dcbb68d10d7a12823 Mon Sep 17 00:00:00 2001 From: Michael Makukha Date: Sat, 4 Jan 2025 01:44:22 +0200 Subject: [PATCH 6/6] Extend test case --- 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 49a53b32..55fdae91 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -20,7 +20,7 @@ def setUp(self): self.tree2 = Tree('a', [ Tree('b', [Token('T', 'x')]), Tree('c', [Token('T', 'y')]), - Tree('d', [Tree('z', [Token('T', 'zz')])]), + Tree('d', [Tree('z', [Token('T', 'zz'), Tree('zzz', 'zzz')])]), ]) def test_eq(self):