Small fix, comments, added some types

Author: erezsh

lark/lark.py

@@ -273,15 +273,15 @@ class Lark(Serialize):
     parser: 'ParsingFrontend'
     terminals: Collection[TerminalDef]
 
-    __serialize_fields__ = 'parser', 'rules', 'options'
+    __serialize_fields__ = ['parser', 'rules', 'options']
 
     def __init__(self, grammar: 'Union[Grammar, str, IO[str]]', **options) -> None:
         self.options = LarkOptions(options)
         re_module: types.ModuleType
 
         # Update which fields are serialized
         if self.options.cache_grammar:
-            self.__serialize_fields__ = self.__serialize_fields__ + ('grammar',)
+            self.__serialize_fields__ = self.__serialize_fields__ + ['grammar']
 
         # Set regex or re module
         use_regex = self.options.regex
@@ -414,7 +414,7 @@ def __init__(self, grammar: 'Union[Grammar, str, IO[str]]', **options) -> None:
             raise ConfigurationError("invalid ambiguity option: %r. Must be one of %r" % (self.options.ambiguity, _VALID_AMBIGUITY_OPTIONS))
 
         if self.options.parser is None:
-            terminals_to_keep = '*'
+            terminals_to_keep = '*'     # For lexer-only mode, keep all terminals
         elif self.options.postlex is not None:
             terminals_to_keep = set(self.options.postlex.always_accept)
         else:

lark/tree_matcher.py

@@ -1,9 +1,10 @@
 """Tree matcher based on Lark grammar"""
 
 import re
+from typing import List, Dict
 from collections import defaultdict
 
-from . import Tree, Token
+from . import Tree, Token, Lark
 from .common import ParserConf
 from .exceptions import ConfigurationError
 from .parsers import earley
@@ -40,7 +41,7 @@ def _best_from_group(seq, group_key, cmp_key):
     return list(d.values())
 
 
-def _best_rules_from_group(rules):
+def _best_rules_from_group(rules: List[Rule]) -> List[Rule]:
     rules = _best_from_group(rules, lambda r: r, lambda r: -len(r.expansion))
     rules.sort(key=lambda r: len(r.expansion))
     return rules
@@ -86,19 +87,23 @@ class TreeMatcher:
 
     Initialize with an instance of Lark.
     """
+    rules_for_root: Dict[str, List[Rule]]
+    rules: List[Rule]
+    parser: Lark
 
-    def __init__(self, parser):
+    def __init__(self, parser: Lark):
         # XXX TODO calling compile twice returns different results!
         assert not parser.options.maybe_placeholders
 
-        # XXX TODO: we just ignore the potential existence of a postlexer
-        if parser.options.postlex is None:
-            self.tokens = parser.terminals.copy()
-            rules = parser.rules.copy()
-        else:
-            if not hasattr(parser, 'grammar') and parser.options.cache:
-                raise ConfigurationError('Unanalyzed grammar not available from cached parser, use cache_grammar=True')
+        if parser.options.postlex and parser.options.postlex.always_accept:
+            # If postlexer's always_accept is used, we need to recompile the grammar with empty terminals-to-keep
+            if not hasattr(parser, 'grammar'):
+                raise ConfigurationError('Source grammar not available from cached parser, use cache_grammar=True'
+                                         if parser.options.cache else "Source grammar not available!")
             self.tokens, rules, _extra = parser.grammar.compile(parser.options.start, set())
+        else:
+            self.tokens = list(parser.terminals)
+            rules = list(parser.rules)
 
         self.rules_for_root = defaultdict(list)
 
@@ -109,9 +114,9 @@ def __init__(self, parser):
         self.rules = _best_rules_from_group(self.rules)
 
         self.parser = parser
-        self._parser_cache = {}
+        self._parser_cache: Dict[str, earley.Parser] = {}
 
-    def _build_recons_rules(self, rules):
+    def _build_recons_rules(self, rules: List[Rule]):
         "Convert tree-parsing/construction rules to tree-matching rules"
         expand1s = {r.origin for r in rules if r.options.expand1}
 
@@ -153,7 +158,7 @@ def _build_recons_rules(self, rules):
                 yield make_recons_rule_to_term(origin, NonTerminal(alias))
             yield make_recons_rule_to_term(origin, origin)
 
-    def match_tree(self, tree, rulename):
+    def match_tree(self, tree: Tree, rulename: str) -> Tree:
         """Match the elements of `tree` to the symbols of rule `rulename`.
 
         Parameters:
@@ -167,7 +172,7 @@ def match_tree(self, tree, rulename):
             UnexpectedToken: If no match was found.
 
         Note:
-            It's the callers' responsibility match the tree recursively.
+            It's the callers' responsibility to match the tree recursively.
         """
         if rulename:
             # validate
@@ -184,11 +189,11 @@ def match_tree(self, tree, rulename):
 
             # TODO pass callbacks through dict, instead of alias?
             callbacks = {rule: rule.alias for rule in rules}
-            conf = ParserConf(rules, callbacks, [rulename])
+            conf = ParserConf(rules, callbacks, [rulename]) # type: ignore[arg-type]
             parser = earley.Parser(self.parser.lexer_conf, conf, _match, resolve_ambiguity=True)
             self._parser_cache[rulename] = parser
 
         # find a full derivation
-        unreduced_tree = parser.parse(ChildrenLexer(tree.children), rulename)
+        unreduced_tree: Tree = parser.parse(ChildrenLexer(tree.children), rulename)
         assert unreduced_tree.data == rulename
         return unreduced_tree