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"""