From e4d73526d4eec4f049fd7737e6f3f813c82aa357 Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Tue, 6 Oct 2020 18:19:52 +0300
Subject: [PATCH] Refactor

---
 lark/lark.py         |  8 +++++++-
 lark/load_grammar.py | 47 ++++++++++++++++++++++----------------------
 2 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/lark/lark.py b/lark/lark.py
index 835a6eb..770b821 100644
--- a/lark/lark.py
+++ b/lark/lark.py
@@ -335,7 +335,13 @@ class Lark(Serialize):
         self._callbacks = None
         # we don't need these callbacks if we aren't building a tree
         if self.options.ambiguity != 'forest':
-            self._parse_tree_builder = ParseTreeBuilder(self.rules, self.options.tree_class or Tree, self.options.propagate_positions, self.options.parser!='lalr' and self.options.ambiguity=='explicit', self.options.maybe_placeholders)
+            self._parse_tree_builder = ParseTreeBuilder(
+                    self.rules,
+                    self.options.tree_class or Tree,
+                    self.options.propagate_positions,
+                    self.options.parser!='lalr' and self.options.ambiguity=='explicit',
+                    self.options.maybe_placeholders
+                )
             self._callbacks = self._parse_tree_builder.create_callback(self.options.transformer)
 
     def _build_parser(self):
diff --git a/lark/load_grammar.py b/lark/load_grammar.py
index 9fb49f6..d039638 100644
--- a/lark/load_grammar.py
+++ b/lark/load_grammar.py
@@ -650,22 +650,6 @@ class Grammar:
 
 
 _imported_grammars = {}
-def import_grammar(grammar_path, loader, base_paths=[]):
-    if grammar_path not in _imported_grammars:
-        import_paths = base_paths + IMPORT_PATHS
-        for import_path in import_paths:
-            with suppress(IOError):
-                joined_path = os.path.join(import_path, grammar_path)
-                with open(joined_path, encoding='utf8') as f:
-                    text = f.read()
-                grammar = loader.load_grammar(text, joined_path)
-                _imported_grammars[grammar_path] = grammar
-                break
-        else:
-            open(grammar_path, encoding='utf8')
-            assert False
-
-    return _imported_grammars[grammar_path]
 
 def import_from_grammar_into_namespace(grammar, namespace, aliases):
     """Returns all rules and terminals of grammar, prepended
@@ -803,7 +787,7 @@ class GrammarLoader:
         ('%ignore expects a value', ['%ignore %import\n']),
     ]
 
-    def __init__(self, re_module, always_keep_all_tokens):
+    def __init__(self, re_module, global_keep_all_tokens):
         terminals = [TerminalDef(name, PatternRE(value)) for name, value in TERMINALS.items()]
 
         rules = [options_from_rule(name, None, x) for name, x in  RULES.items()]
@@ -816,7 +800,24 @@ class GrammarLoader:
 
         self.canonize_tree = CanonizeTree()
         self.re_module = re_module
-        self.always_keep_all_tokens = always_keep_all_tokens
+        self.global_keep_all_tokens = global_keep_all_tokens
+
+    def import_grammar(self, grammar_path, base_paths=[]):
+        if grammar_path not in _imported_grammars:
+            import_paths = base_paths + IMPORT_PATHS
+            for import_path in import_paths:
+                with suppress(IOError):
+                    joined_path = os.path.join(import_path, grammar_path)
+                    with open(joined_path, encoding='utf8') as f:
+                        text = f.read()
+                    grammar = self.load_grammar(text, joined_path)
+                    _imported_grammars[grammar_path] = grammar
+                    break
+            else:
+                open(grammar_path, encoding='utf8') # Force a file not found error
+                assert False
+
+        return _imported_grammars[grammar_path]
 
     def load_grammar(self, grammar_text, grammar_name='<?>'):
         "Parse grammar_text, verify, and create Grammar object. Display nice messages on error."
@@ -902,7 +903,7 @@ class GrammarLoader:
         # import grammars
         for dotted_path, (base_paths, aliases) in imports.items():
             grammar_path = os.path.join(*dotted_path) + EXT
-            g = import_grammar(grammar_path, self, base_paths=base_paths)
+            g = self.import_grammar(grammar_path, base_paths=base_paths)
             new_td, new_rd = import_from_grammar_into_namespace(g, '__'.join(dotted_path), aliases)
 
             term_defs += new_td
@@ -949,9 +950,9 @@ class GrammarLoader:
         rule_names = {}
         for name, params, _x, option in rules:
             # We can't just simply not throw away the tokens later, we need option.keep_all_tokens to correctly generate maybe_placeholders
-            if self.always_keep_all_tokens: 
+            if self.global_keep_all_tokens:
                 option.keep_all_tokens = True
-            
+
             if name.startswith('__'):
                 raise GrammarError('Names starting with double-underscore are reserved (Error at %s)' % name)
             if name in rule_names:
@@ -986,5 +987,5 @@ class GrammarLoader:
 
 
 
-def load_grammar(grammar, source, re_, always_keep_all_tokens):
-    return GrammarLoader(re_, always_keep_all_tokens).load_grammar(grammar, source)
+def load_grammar(grammar, source, re_, global_keep_all_tokens):
+    return GrammarLoader(re_, global_keep_all_tokens).load_grammar(grammar, source)