Add template support to Reconstructor

lark/reconstruct.py

@@ -7,9 +7,9 @@
 from .tree import Tree, ParseTree
 from .visitors import Transformer_InPlace
 from .lexer import Token, PatternStr, TerminalDef
-from .grammar import Terminal, NonTerminal, Symbol
+from .grammar import Terminal, Symbol
 
-from .tree_matcher import TreeMatcher, is_discarded_terminal
+from .tree_matcher import TreeMatcher, is_discarded_terminal, parse_rulename
 from .utils import is_id_continue
 
 def is_iter_empty(i):
@@ -55,7 +55,9 @@ def __default__(self, data, children, meta):
                     if isinstance(x, Token):
                         assert Terminal(x.type) == sym, x
                     else:
-                        assert NonTerminal(x.data) == sym, (sym, x)
+                        name, _args = parse_rulename(sym.name)
+                        assert x.data == name, (sym, x)
+                        x.meta.original_rulename = sym.name
                     to_write.append(x)
 
         assert is_iter_empty(iter_args)
@@ -82,14 +84,14 @@ def __init__(self, parser: Lark, term_subs: Optional[Dict[str, Callable[[Symbol]
 
         self.write_tokens = WriteTokensTransformer({t.name:t for t in self.tokens}, term_subs or {})
 
-    def _reconstruct(self, tree):
-        unreduced_tree = self.match_tree(tree, tree.data)
+    def _reconstruct(self, tree, original_rulename=None):
+        unreduced_tree = self.match_tree(tree, original_rulename or tree.data)
 
         res = self.write_tokens.transform(unreduced_tree)
         for item in res:
             if isinstance(item, Tree):
-                # TODO use orig_expansion.rulename to support templates
-                yield from self._reconstruct(item)
+                child_rulename = getattr(item.meta, "original_rulename", None)
+                yield from self._reconstruct(item, child_rulename)
             else:
                 yield item
 

tests/test_reconstructor.py

@@ -101,6 +101,20 @@ def test_expand_rule(self):
         for c in code:
             self.assert_reconstruct(g, c)
 
+    def test_template_rule(self):
+        g = """
+        start: (NL | binfunc)*
+        binfunc: WORD "(" binarg{WORD} ")"
+        binarg{arg}: arg "," arg
+        NL: /(\\r?\\n)+\\s*/
+        """ + common
+
+        code = """
+        func(a, b)
+        """
+
+        self.assert_reconstruct(g, code)
+
     def test_json_example(self):
         test_json = '''
             {