Fixes caching when custom lexers are used

lark-stubs/lark.pyi

@@ -55,6 +55,7 @@ class Lark:
         propagate_positions: bool = False,
         maybe_placeholders: bool = False,
         lexer_callbacks: Optional[Dict[str, Callable[[Token], Token]]] = None,
+        cache: Union[bool, str] = False,
         g_regex_flags: int = ...
     ):
         ...

lark/parser_frontends.py

@@ -1,5 +1,3 @@
-from functools import partial
-
 from .utils import get_regexp_width, Serialize
 from .parsers.grammar_analysis import GrammarAnalyzer
 from .lexer import TraditionalLexer, ContextualLexer, Lexer, Token
@@ -20,7 +18,14 @@ def get_frontend(parser, lexer):
         elif lexer == 'contextual':
             return LALR_ContextualLexer
         elif issubclass(lexer, Lexer):
-            return partial(LALR_CustomLexer, lexer)
+            class LALR_CustomLexerWrapper(LALR_CustomLexer):
+                def __init__(self, lexer_conf, parser_conf, options=None):
+                    super(LALR_CustomLexerWrapper, self).__init__(
+                        lexer, lexer_conf, parser_conf, options=options)
+                def init_lexer(self):
+                    self.lexer = lexer(self.lexer_conf)
+
+            return LALR_CustomLexerWrapper
         else:
             raise ValueError('Unknown lexer: %s' % lexer)
     elif parser=='earley':

tests/test_cache.py

@@ -4,6 +4,7 @@ import sys
 from unittest import TestCase, main
 
 from lark import Lark, Tree
+from lark.lexer import Lexer, Token
 import lark.lark as lark_module
 
 try:
@@ -38,6 +39,15 @@ class MockFS:
         return name in self.files
 
 
+class CustomLexer(Lexer):
+    def __init__(self, lexer_conf):
+        pass
+
+    def lex(self, data):
+        for obj in data:
+            yield Token('A', obj)
+
+
 class TestCache(TestCase):
     def setUp(self):
         pass
@@ -70,6 +80,12 @@ class TestCache(TestCase):
             parser = Lark(g, parser='lalr', cache=True)
             assert parser.parse('a') == Tree('start', [])
 
+            # Test with custom lexer
+            mock_fs.files = {}
+            parser = Lark(g, parser='lalr', lexer=CustomLexer, cache=True)
+            parser = Lark(g, parser='lalr', lexer=CustomLexer, cache=True)
+            assert len(mock_fs.files) == 1
+            assert parser.parse('a') == Tree('start', [])
         finally:
             lark_module.FS = fs