tuple -> Tuple

Author: MegaIng

lark/load_grammar.py

@@ -101,13 +101,14 @@
     '_DECLARE': r'%declare',
     '_EXTEND': r'%extend',
     '_IMPORT': r'%import',
+    '_INCLUDE': r'%include',
     'NUMBER': r'[+-]?\d+',
 }
 
 RULES = {
     'start': ['_list'],
     '_list':  ['_item', '_list _item'],
-    '_item':  ['rule', 'term', 'ignore', 'import', 'declare', 'override', 'extend', '_NL'],
+    '_item':  ['rule', 'term', 'ignore', 'import', 'declare', 'override', 'extend', 'include', '_NL'],
 
     'rule': ['RULE template_params _COLON expansions _NL',
              'RULE template_params _DOT NUMBER _COLON expansions _NL'],
@@ -164,6 +165,7 @@
     'import': ['_IMPORT _import_path _NL',
                '_IMPORT _import_path _LPAR name_list _RPAR _NL',
                '_IMPORT _import_path _TO name _NL'],
+    'include': ['_INCLUDE _import_path _NL'],
 
     '_import_path': ['import_lib', 'import_rel'],
     'import_lib': ['_import_args'],
@@ -811,7 +813,7 @@ def __init__(self, pkg_name: str, search_paths: Tuple[str, ...]=("", )) -> None:
     def __repr__(self):
         return "%s(%r, %r)" % (type(self).__name__, self.pkg_name, self.search_paths)
 
-    def __call__(self, base_path: Union[None, str, PackageResource], grammar_path: str, used_files: Dict[str, tuple[str, str]]=None) -> Tuple[PackageResource, str]:
+    def __call__(self, base_path: Union[None, str, PackageResource], grammar_path: str, used_files: Dict[str, Tuple[str, str]]=None) -> Tuple[PackageResource, str]:
         if base_path is None:
             to_try = self.search_paths
         else:
@@ -1137,7 +1139,10 @@ def _unpack_import(self, stmt, grammar_name):
             path_node, = stmt.children
             arg1 = None
 
-        if isinstance(arg1, Tree):  # Multi import
+        if stmt.data == "include": # we only want the dotted_path, no aliases
+            dotted_path = tuple(path_node.children)
+            aliases = None
+        elif isinstance(arg1, Tree):  # Multi import
             dotted_path = tuple(path_node.children)
             names = arg1.children
             aliases = dict(zip(names, names))  # Can't have aliased multi import, so all aliases will be the same as names
@@ -1185,50 +1190,66 @@ def _unpack_definition(self, tree, mangle):
         exp = _mangle_exp(exp, mangle)
         return name, exp, params, opts
 
-
-    def load_grammar(self, grammar_text: str, grammar_name: str="<?>", mangle: Optional[Callable[[str], str]]=None) -> None:
+    def _get_parsed(self, grammar_text, grammar_name):
         if grammar_text not in self.cached_grammars:
             tree = _parse_grammar(grammar_text, grammar_name)
             self.cached_grammars[grammar_text] = tree
-        tree = nr_deepcopy_tree(self.cached_grammars[grammar_text])
+        return nr_deepcopy_tree(self.cached_grammars[grammar_text])
+
+    def load_grammar(self, grammar_text, grammar_name="<?>", mangle=None):
+        tree = self._get_parsed(grammar_text, grammar_name)
 
         imports = {}
-        for stmt in tree.children:
-            if stmt.data == 'import':
-                dotted_path, base_path, aliases = self._unpack_import(stmt, grammar_name)
-                try:
-                    import_base_path, import_aliases = imports[dotted_path]
-                    assert base_path == import_base_path, 'Inconsistent base_path for %s.' % '.'.join(dotted_path)
-                    import_aliases.update(aliases)
-                except KeyError:
-                    imports[dotted_path] = base_path, aliases
+        todo = [tree.children]
+        after_import_todo = []
+        while todo:
+            stmts = todo.pop(0)
+            after_import_todo.append(stmts)
+            for stmt in stmts:
+                if stmt.data == 'import':
+                    dotted_path, base_path, aliases = self._unpack_import(stmt, grammar_name)
+                    try:
+                        import_base_path, import_aliases = imports[dotted_path]
+                        assert base_path == import_base_path, 'Inconsistent base_path for %s.' % '.'.join(dotted_path)
+                        import_aliases.update(aliases)
+                    except KeyError:
+                        imports[dotted_path] = base_path, aliases
+                elif stmt.data == 'include':
+                    dotted_path, base_path, aliases = self._unpack_import(stmt, grammar_name)
+                    assert aliases is None, "For some reason '_unpack_import' returned aliases for %include"
+                    text, name = self.resolve_import(dotted_path, base_path)
+                    new_tree = self._get_parsed(text, name)
+                    todo.append(new_tree.children)
+
 
         for dotted_path, (base_path, aliases) in imports.items():
             self.do_import(dotted_path, base_path, aliases, mangle)
 
-        for stmt in tree.children:
-            if stmt.data in ('term', 'rule'):
-                self._define(*self._unpack_definition(stmt, mangle))
-            elif stmt.data == 'override':
-                r ,= stmt.children
-                self._define(*self._unpack_definition(r, mangle), override=True)
-            elif stmt.data == 'extend':
-                r ,= stmt.children
-                self._extend(*self._unpack_definition(r, mangle))
-            elif stmt.data == 'ignore':
-                # if mangle is not None, we shouldn't apply ignore, since we aren't in a toplevel grammar
-                if mangle is None:
-                    self._ignore(*stmt.children)
-            elif stmt.data == 'declare':
-                names = [t.value for t in stmt.children]
-                if mangle is None:
-                    self._declare(*names)
+        while after_import_todo:
+            stmts = after_import_todo.pop(0)
+            for stmt in stmts:
+                if stmt.data in ('term', 'rule'):
+                    self._define(*self._unpack_definition(stmt, mangle))
+                elif stmt.data == 'override':
+                    r ,= stmt.children
+                    self._define(*self._unpack_definition(r, mangle), override=True)
+                elif stmt.data == 'extend':
+                    r ,= stmt.children
+                    self._extend(*self._unpack_definition(r, mangle))
+                elif stmt.data == 'ignore':
+                    # if mangle is not None, we shouldn't apply ignore, since we aren't in a toplevel grammar
+                    if mangle is None:
+                        self._ignore(*stmt.children)
+                elif stmt.data == 'declare':
+                    names = [t.value for t in stmt.children]
+                    if mangle is None:
+                        self._declare(*names)
+                    else:
+                        self._declare(*map(mangle, names))
+                elif stmt.data in ('import', 'include'): # These have already been processed
+                    pass
                 else:
-                    self._declare(*map(mangle, names))
-            elif stmt.data == 'import':
-                pass
-            else:
-                assert False, stmt
+                    assert False, stmt
 
 
         term_defs = { name: exp
@@ -1252,9 +1273,8 @@ def rule_dependencies(symbol):
         self._definitions = {k: v for k, v in self._definitions.items() if k in _used}
 
 
-    def do_import(self, dotted_path: Tuple[str, ...], base_path: Optional[str], aliases: Dict[str, str], base_mangle: Optional[Callable[[str], str]]=None) -> None:
+    def resolve_import(self, dotted_path: Tuple[str, ...], base_path: Optional[str]) -> Tuple[str, Union[str, PackageResource]]:
         assert dotted_path
-        mangle = _get_mangle('__'.join(dotted_path), aliases, base_mangle)
         grammar_path = os.path.join(*dotted_path) + EXT
         to_try = self.import_paths + ([base_path] if base_path is not None else []) + [stdlib_loader]
         for source in to_try:
@@ -1274,21 +1294,26 @@ def do_import(self, dotted_path: Tuple[str, ...], base_path: Optional[str], alia
                 h = hashlib.md5(text.encode('utf8')).hexdigest()
                 if self.used_files.setdefault(joined_path, (h,text))[0] != h:
                     raise RuntimeError("Grammar file was changed during importing")
-
-                gb = GrammarBuilder(self.global_keep_all_tokens, self.import_paths, self.used_files, self.cached_grammars)
-                gb.load_grammar(text, joined_path, mangle)
-                gb._remove_unused(map(mangle, aliases))
-                for name in gb._definitions:
-                    if name in self._definitions:
-                        raise GrammarError("Cannot import '%s' from '%s': Symbol already defined." % (name, grammar_path))
-
-                self._definitions.update(**gb._definitions)
-                break
+                return text, joined_path
         else:
             # Search failed. Make Python throw a nice error.
             open(grammar_path, encoding='utf8')
-            assert False, "Couldn't import grammar %s, but a corresponding file was found at a place where lark doesn't search for it" % (dotted_path,)
+            raise GrammarError("Couldn't import grammar %s, but a corresponding file was found at a place where lark"
+                               " doesn't search for it. This probably means that you wrongly used non-relative imports."
+                               " Try adding a `.` at the beginning of the path" % (dotted_path,))
+
+    def do_import(self, dotted_path: Tuple[str, ...], base_path: Optional[str], aliases: Dict[str, str], base_mangle: Optional[Callable[[str], str]]=None):
+        mangle = _get_mangle('__'.join(dotted_path), aliases, base_mangle)
+        text, joined_path = self.resolve_import(dotted_path, base_path)
+
+        gb = GrammarBuilder(self.global_keep_all_tokens, self.import_paths, self.used_files)
+        gb.load_grammar(text, joined_path, mangle)
+        gb._remove_unused(map(mangle, aliases))
+        for name in gb._definitions:
+            if name in self._definitions:
+                raise GrammarError("Cannot import '%s' from '%s': Symbol already defined." % (name, joined_path))
 
+        self._definitions.update(**gb._definitions)
 
     def validate(self) -> None:
         for name, (params, exp, _options) in self._definitions.items():