From 2f4831f9b6dd857dcb3b8d53a8839474d3c5e5f7 Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Thu, 13 Aug 2020 21:13:42 +0300
Subject: [PATCH] Small refactor after PR

---
 lark/lark.py             | 12 +++++-------
 lark/parser_frontends.py |  9 ++++++---
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/lark/lark.py b/lark/lark.py
index 3ed96d7..8371943 100644
--- a/lark/lark.py
+++ b/lark/lark.py
@@ -11,7 +11,7 @@ from .common import LexerConf, ParserConf
 
 from .lexer import Lexer, TraditionalLexer, TerminalDef, UnexpectedToken
 from .parse_tree_builder import ParseTreeBuilder
-from .parser_frontends import get_frontend
+from .parser_frontends import get_frontend, _get_lexer_callbacks
 from .grammar import Rule
 
 import re
@@ -278,12 +278,10 @@ class Lark(Serialize):
                     rule.options.priority = None
 
         # TODO Deprecate lexer_callbacks?
-        lexer_callbacks = dict(self.options.lexer_callbacks)
-        if self.options.transformer:
-            t = self.options.transformer
-            for term in self.terminals:
-                if hasattr(t, term.name):
-                    lexer_callbacks[term.name] = getattr(t, term.name)
+        lexer_callbacks = (_get_lexer_callbacks(self.options.transformer, self.terminals)
+                           if self.options.transformer
+                           else {})
+        lexer_callbacks.update(self.options.lexer_callbacks)
 
         self.lexer_conf = LexerConf(self.terminals, re_module, self.ignore_tokens, self.options.postlex, lexer_callbacks, self.options.g_regex_flags, use_bytes=self.options.use_bytes)
 
diff --git a/lark/parser_frontends.py b/lark/parser_frontends.py
index a45bf9c..b993b9f 100644
--- a/lark/parser_frontends.py
+++ b/lark/parser_frontends.py
@@ -58,9 +58,8 @@ class _ParserFrontend(Serialize):
         return self.parser.parse(input, start, *args)
 
 
-def _recreate_lexer_callbacks(memo, transformer):
+def _get_lexer_callbacks(transformer, terminals):
     result = {}
-    terminals = [item for item in memo.values() if isinstance(item, TerminalDef)]
     for terminal in terminals:
         callback = getattr(transformer, terminal.name, None)
         if callback is not None:
@@ -85,12 +84,16 @@ class WithLexer(_ParserFrontend):
     @classmethod
     def deserialize(cls, data, memo, callbacks, postlex, transformer, re_module):
         inst = super(WithLexer, cls).deserialize(data, memo)
+
         inst.postlex = postlex
         inst.parser = LALR_Parser.deserialize(inst.parser, memo, callbacks)
-        inst.lexer_conf.callbacks = _recreate_lexer_callbacks(memo, transformer)
+
+        terminals = [item for item in memo.values() if isinstance(item, TerminalDef)]
+        inst.lexer_conf.callbacks = _get_lexer_callbacks(transformer, terminals)
         inst.lexer_conf.re_module = re_module
         inst.lexer_conf.skip_validation=True
         inst.init_lexer()
+
         return inst
 
     def _serialize(self, data, memo):