From 5930e4ba6ff28df59aa2cdc0196373a1c37d1970 Mon Sep 17 00:00:00 2001
From: MegaIng1 <trampchamp@hotmail.de>
Date: Fri, 25 Dec 2020 01:26:41 +0100
Subject: [PATCH] Added Terminal support for `%override`

---
 docs/grammar.md       |  2 +-
 lark/load_grammar.py  | 25 +++++++++++++++++++++----
 tests/test_grammar.py | 18 +++++++++++++++---
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/docs/grammar.md b/docs/grammar.md
index b899b3f..d6d4b3b 100644
--- a/docs/grammar.md
+++ b/docs/grammar.md
@@ -291,7 +291,7 @@ Declare a terminal without defining it. Useful for plugins.
 
 ### %override
 
-Override a rule, affecting all the rules that refer to it.
+Override a rule or terminals, affecting all references to it, even in imported grammars.
 
 Useful for implementing an inheritance pattern when importing grammars.
 
diff --git a/lark/load_grammar.py b/lark/load_grammar.py
index 0fafc1c..7383c17 100644
--- a/lark/load_grammar.py
+++ b/lark/load_grammar.py
@@ -149,8 +149,8 @@ RULES = {
 
     'term': ['TERMINAL _COLON expansions _NL',
              'TERMINAL _DOT NUMBER _COLON expansions _NL'],
-    'statement': ['ignore', 'import', 'declare', 'override_rule'],
-    'override_rule': ['_OVERRIDE rule'],
+    'statement': ['ignore', 'import', 'declare', 'override'],
+    'override': ['_OVERRIDE rule', '_OVERRIDE term'],
     'ignore': ['_IGNORE expansions _NL'],
     'declare': ['_DECLARE _declare_args _NL'],
     'import': ['_IMPORT _import_path _NL',
@@ -950,6 +950,7 @@ class GrammarLoader:
         # Execute statements
         ignore, imports = [], {}
         overriding_rules = []
+        overriding_terms = []
         for (stmt,) in statements:
             if stmt.data == 'ignore':
                 t ,= stmt.children
@@ -998,9 +999,15 @@ class GrammarLoader:
             elif stmt.data == 'declare':
                 for t in stmt.children:
                     term_defs.append([t.value, (None, None)])
-            elif stmt.data == 'override_rule':
+            elif stmt.data == 'override':
                 r ,= stmt.children
-                overriding_rules.append(options_from_rule(*r.children))
+                if r.data == 'rule':
+                    overriding_rules.append(options_from_rule(*r.children))
+                else:
+                    if len(r.children) == 2:
+                        overriding_terms.append((r.children[0].value, (r.children[1], 1)))
+                    else:
+                        overriding_terms.append((r.children[0].value, (r.children[2], int(r.children[1]))))
             else:
                 assert False, stmt
 
@@ -1022,6 +1029,16 @@ class GrammarLoader:
                 raise GrammarError("Cannot override a nonexisting rule: %s" % name)
             rule_defs.append(r)
 
+        # Same for terminals
+        for t in overriding_terms:
+            name = t[0]
+            # remove overridden rule from rule_defs
+            overridden, term_defs = classify_bool(term_defs, lambda t: t[0] == name)    # FIXME inefficient
+            if not overridden:
+                raise GrammarError("Cannot override a nonexisting terminal: %s" % name)
+            term_defs.append(t)
+
+
         ## Handle terminals
 
         # Verify correctness 1
diff --git a/tests/test_grammar.py b/tests/test_grammar.py
index 3ce76f6..760d563 100644
--- a/tests/test_grammar.py
+++ b/tests/test_grammar.py
@@ -3,7 +3,7 @@ from __future__ import absolute_import
 import sys
 from unittest import TestCase, main
 
-from lark import Lark
+from lark import Lark, Token
 from lark.load_grammar import GrammarLoader, GrammarError
 
 
@@ -21,7 +21,7 @@ class TestGrammar(TestCase):
                 else:
                     assert False, "example did not raise an error"
 
-    def test_override(self):
+    def test_override_rule(self):
         # Overrides the 'sep' template in existing grammar to add an optional terminating delimiter
         # Thus extending it beyond its original capacity
         p = Lark("""
@@ -29,12 +29,24 @@ class TestGrammar(TestCase):
 
             %override sep{item, delim}: item (delim item)* delim?
             %ignore " "
-        """)
+        """, source_path=__file__)
 
         a = p.parse('[1, 2, 3]')
         b = p.parse('[1, 2, 3, ]')
         assert a == b
 
+    def test_override_terminal(self):
+        p = Lark("""
+            
+            %import .grammars.ab (startab, A, B)
+            
+            %override A: "C"
+            %override B: "D"
+        """, start='startab', source_path=__file__)
+
+        a = p.parse('CD')
+        self.assertEqual(a.children[0].children, [Token('A', 'C'), Token('B', 'D')])
+
 
 
 if __name__ == '__main__':