From b3816b9a90abcf914f12891cd6a6fbb74f0ee771 Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Fri, 27 Aug 2021 12:59:31 +0100
Subject: [PATCH] Added utility function: List grammar imports + test

---
 lark/load_grammar.py  |  6 ++++++
 tests/test_grammar.py | 18 ++++++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/lark/load_grammar.py b/lark/load_grammar.py
index 7c64196..4a9360c 100644
--- a/lark/load_grammar.py
+++ b/lark/load_grammar.py
@@ -1336,6 +1336,12 @@ def verify_used_files(file_hashes):
             return False
     return True
 
+def list_grammar_imports(grammar, import_paths=[]):
+    "Returns a list of paths to the lark grammars imported by the given grammar (recursively)"
+    builder = GrammarBuilder(False, import_paths)
+    builder.load_grammar(grammar, '<string>')
+    return list(builder.used_files.keys())
+
 def load_grammar(grammar, source, import_paths, global_keep_all_tokens):
     builder = GrammarBuilder(global_keep_all_tokens, import_paths)
     builder.load_grammar(grammar, source)
diff --git a/tests/test_grammar.py b/tests/test_grammar.py
index 47a345c..319d709 100644
--- a/tests/test_grammar.py
+++ b/tests/test_grammar.py
@@ -1,10 +1,10 @@
 from __future__ import absolute_import
 
-import sys
+import os
 from unittest import TestCase, main
 
 from lark import Lark, Token, Tree, ParseError, UnexpectedInput
-from lark.load_grammar import GrammarError, GRAMMAR_ERRORS, find_grammar_errors
+from lark.load_grammar import GrammarError, GRAMMAR_ERRORS, find_grammar_errors, list_grammar_imports
 from lark.load_grammar import FromPackageLoader
 
 
@@ -258,6 +258,20 @@ class TestGrammar(TestCase):
         for i in (-1, 1000):
             self.assertRaises(UnexpectedInput, l.parse, str(i))
 
+    def test_list_grammar_imports(self):
+            grammar = """
+            %import .test_templates_import (start, sep)
+
+            %override sep{item, delim}: item (delim item)* delim?
+            %ignore " "
+            """
+
+            imports = list_grammar_imports(grammar, [os.path.dirname(__file__)])
+            self.assertEqual({os.path.split(i)[-1] for i in imports}, {'test_templates_import.lark', 'templates.lark'})
+
+            imports = list_grammar_imports('%import common.WS', [])
+            assert len(imports) == 1 and imports[0].pkg_name == 'lark'
+
 
 if __name__ == '__main__':
     main()