diff --git a/lark/reconstruct.py b/lark/reconstruct.py index 2d8423aec..2bb7221db 100644 --- a/lark/reconstruct.py +++ b/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 diff --git a/tests/test_reconstructor.py b/tests/test_reconstructor.py index a3cdcbd80..d26ccbb3d 100644 --- a/tests/test_reconstructor.py +++ b/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 = ''' {