Repository: lark-parser/lark
Pull request #1096 (2022-02-05 22:43:18)

Add test cases for tree templates
Also include minor fixes for clear issues

The test cases probably aren't the best. The were mostly written by going backwards for code coverage instead of starting with a strong understanding of the intended API.

While I was working on this, I did notice that most methods are annotated to take in `TreeOrCode`. `_get_tree()` is then used to parse the value if the argument is a `str`. However, `TemplateConf` isn't guaranteed to have a parse function. I guess the idea is that if you already working with just `Tree`s, you don't need one. But the current API makes it very easy to get an `AssertionError` if the classes are not create properly & type annotations aren't going to be able to protect against that.

erezsh: Looks good, thanks!

Patch:
diff --git a/lark/exceptions.py b/lark/exceptions.py
index 5c65bb9..a3d7d90 100644
--- a/lark/exceptions.py
+++ b/lark/exceptions.py
@@ -279,4 +279,8 @@ class VisitError(LarkError):
         self.obj = obj
         self.orig_exc = orig_exc
 
+
+class MissingVariableError(LarkError):
+    pass
+
 ###}
diff --git a/lark/tree_templates.py b/lark/tree_templates.py
index 8db8b35..78916ad 100644
--- a/lark/tree_templates.py
+++ b/lark/tree_templates.py
@@ -7,9 +7,11 @@ A tree templates is a tree that contains nodes that are template variables.
 from typing import Union, Optional, Mapping
 
 from lark import Tree, Transformer
+from lark.exceptions import MissingVariableError
 
 TreeOrCode = Union[Tree, str]
 
+
 class TemplateConf:
     """Template Configuration
 
@@ -19,7 +21,6 @@ class TemplateConf:
     def __init__(self, parse=None):
         self._parse = parse
 
-
     def test_var(self, var: Union[Tree, str]) -> Optional[str]:
         """Given a tree node, if it is a template variable return its name. Otherwise, return None.
 
@@ -32,12 +33,13 @@ class TemplateConf:
         if isinstance(var, str) and var.startswith('$'):
             return var.lstrip('$')
 
-        if isinstance(var, Tree) and var.data == 'var' and var.children[0].startswith('$'):
-            return var.children[0].lstrip('$')
+        if isinstance(var, Tree) and var.data == 'var' and len(var.children) > 0:
+            first_child = var.children[0]
+            if isinstance(first_child, str) and first_child.startswith('$'):
+                return first_child.lstrip('$')
 
         return None
 
-
     def _get_tree(self, template: TreeOrCode):
         if isinstance(template, str):
             assert self._parse
@@ -57,7 +59,7 @@ class TemplateConf:
         if isinstance(template, str):
             if template == tree:
                 return {}
-            return
+            return None
 
         assert isinstance(template, Tree), template
 
@@ -66,14 +68,13 @@ class TemplateConf:
             for t1, t2 in zip(template.children, tree.children):
                 matches = self._match_tree_template(t1, t2)
                 if matches is None:
-                    return
+                    return None
 
                 res.update(matches)
 
             return res
 
 
-
 class _ReplaceVars(Transformer):
     def __init__(self, conf, vars):
         self._conf = conf
@@ -84,7 +85,10 @@ class _ReplaceVars(Transformer):
 
         var = self._conf.test_var(tree)
         if var:
-            return self._vars[var]
+            try:
+                return self._vars[var]
+            except KeyError:
+                raise MissingVariableError(f"No mapping for template variable ({var})")
         return tree
 
 
@@ -141,7 +145,6 @@ def translate(t1: Template, t2: Template, tree: TreeOrCode):
     return tree
 
 
-
 class TemplateTranslator:
     """Utility class for translating a collection of patterns
     """


Test patch:
diff --git a/tests/__main__.py b/tests/__main__.py
index ed45960..bbf9986 100644
--- a/tests/__main__.py
+++ b/tests/__main__.py
@@ -12,6 +12,7 @@ from .test_reconstructor import TestReconstructor
 from .test_tree_forest_transformer import TestTreeForestTransformer
 from .test_lexer import TestLexer
 from .test_python_grammar import TestPythonParser
+from .test_tree_templates import *  # We define __all__ to list which TestSuites to run
 
 try:
     from .test_nearley.test_nearley import TestNearley
@@ -23,7 +24,7 @@ except ImportError:
 
 from .test_logger import Testlogger
 
-from .test_parser import * # We define __all__ to list which TestSuites to run
+from .test_parser import *  # We define __all__ to list which TestSuites to run
 
 logger.setLevel(logging.INFO)
 
diff --git a/tests/test_tree_templates.py b/tests/test_tree_templates.py
new file mode 100644
index 0000000..fa9578e
--- /dev/null
+++ b/tests/test_tree_templates.py
@@ -0,0 +1,225 @@
+from __future__ import absolute_import
+
+import unittest
+from copy import deepcopy
+
+from lark import Lark, Tree, Token
+from lark.exceptions import MissingVariableError
+from lark.tree_templates import TemplateConf, Template, TemplateTranslator
+
+SOME_NON_TEMPLATED_STRING = "foo bar"
+SOME_TEMPLATE_NAME = "thing"
+SOME_TEMPLATE_STRING = f"${SOME_TEMPLATE_NAME}"
+SOME_NON_STRING = 12345
+SOME_TEMPLATING_GRAMMAR = r"""
+start: DASHES? foo DASHES? bar
+DASHES: "--"
+foo: "foo"
+    | TEMPLATE_NAME -> var
+bar: "bar"
+    | TEMPLATE_NAME -> var
+TEMPLATE_NAME: "$" NAME
+NAME: /[^\W\d]\w*/
+%ignore /[\t \f]+/  // WS
+"""
+SOME_FOO_TEMPLATE = f"{SOME_TEMPLATE_STRING} bar"
+SOME_BAR_TEMPLATE = f"foo {SOME_TEMPLATE_STRING}"
+SOME_NON_TEMPLATE_TREE = Tree("foo", children=["hi"])
+
+__all__ = [
+    "TestTreeTemplatesConf",
+    "TestTreeTemplatesTemplateTranslator",
+    "TestTreeTemplatesTemplate",
+    "TestTreeTemplatesTemplateDefaultConf",
+]
+
+
+class TestTreeTemplatesConf(unittest.TestCase):
+    parser = Lark(SOME_TEMPLATING_GRAMMAR)
+
+    def test_conf_test_var__not_var(self):
+        conf = TemplateConf(self.parser.parse)
+
+        non_templates = {
+            "non-templated string": SOME_NON_TEMPLATED_STRING,
+            "non-var tree": Tree("stmt", children=[]),
+            "var tree, non-templated string": Tree(
+                "var", children=[SOME_NON_TEMPLATED_STRING]
+            ),
+            "var tree, templated string not first child": Tree(
+                "var", children=[SOME_NON_TEMPLATED_STRING, SOME_TEMPLATE_STRING]
+            ),
+            "var tree, first child not string": Tree("var", children=[SOME_NON_STRING]),
+            "var tree, no children": Tree("var", children=[]),
+        }
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertIsNone(conf.test_var(test_case))
+
+    def test_conf_test_var__is_var(self):
+        conf = TemplateConf(self.parser.parse)
+
+        non_templates = {
+            "templated string": SOME_TEMPLATE_STRING,
+            "var tree, non-templated string": Tree(
+                "var", children=[SOME_TEMPLATE_STRING]
+            ),
+        }
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertEqual(SOME_TEMPLATE_NAME, conf.test_var(test_case))
+
+    def test_conf_call__same_tree(self):
+        conf = TemplateConf(self.parser.parse)
+        explicitly_parsed = self.parser.parse(SOME_FOO_TEMPLATE)
+
+        non_templates = {
+            "to be parsed": SOME_FOO_TEMPLATE,
+            "already parsed": explicitly_parsed,
+        }
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                template = conf(test_case)
+                self.assertEqual(explicitly_parsed, template.tree)
+
+    def test_template_match__default_conf_match_same_tree__empty_dictionary(self):
+        template = Template(SOME_NON_TEMPLATE_TREE)
+
+        self.assertEqual({}, template.match(SOME_NON_TEMPLATE_TREE))
+
+
+class TestTreeTemplatesTemplate(unittest.TestCase):
+    parser = Lark(SOME_TEMPLATING_GRAMMAR)
+    conf = TemplateConf(parser.parse)
+
+    def test_template_match__same_tree_no_template__empty_dictionary(self):
+        template = Template(SOME_NON_TEMPLATE_TREE, conf=self.conf)
+
+        self.assertEqual({}, template.match(SOME_NON_TEMPLATE_TREE))
+
+    def test_template_match__different_tree_no_template__none(self):
+        template = Template(SOME_NON_TEMPLATE_TREE, conf=self.conf)
+
+        self.assertIsNone(template.match(Tree("foo", children=["bye"])))
+
+    def test_template_match__no_template__empty_dictionary(self):
+        tree = self.parser.parse(SOME_NON_TEMPLATED_STRING)
+        template = Template(tree, conf=self.conf)
+
+        non_templates = {
+            "un-parsed string": SOME_NON_TEMPLATED_STRING,
+            "parsed tree": tree,
+        }
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertEqual({}, template.match(test_case))
+
+    def test_template_match__with_template__empty_dictionary(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        non_templates = {"un-parsed string": SOME_FOO_TEMPLATE, "parsed tree": tree}
+        expected_result = {SOME_TEMPLATE_NAME: tree.children[0]}
+
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertEqual(expected_result, template.match(test_case))
+
+    def test_template_match__different_tree__none(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        non_templates = {
+            "un-parsed string": SOME_BAR_TEMPLATE,
+            "parsed tree": self.parser.parse(SOME_BAR_TEMPLATE),
+        }
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertIsNone(template.match(test_case))
+
+    def test_template_search__same_tree_no_template__empty_generator(self):
+        template = Template(SOME_NON_TEMPLATE_TREE, conf=self.conf)
+
+        self.assertEqual([], list(template.search(SOME_NON_TEMPLATE_TREE)))
+
+    def test_template_search__same_tree_as_child__empty_generator(self):
+        template = Template(SOME_NON_TEMPLATE_TREE, conf=self.conf)
+
+        self.assertEqual(
+            [], list(template.search(Tree("root", children=[SOME_NON_TEMPLATE_TREE])))
+        )
+
+    def test_template_search__with_template__matched_result_with_parent_tree(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        non_templates = {"un-parsed string": SOME_FOO_TEMPLATE, "parsed tree": tree}
+        expected_result = [(tree, {SOME_TEMPLATE_NAME: tree.children[0]})]
+
+        for description, test_case in non_templates.items():
+            with self.subTest(msg=description):
+                self.assertEqual(expected_result, list(template.search(test_case)))
+
+    def test_template_apply_vars__empty__exception(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        with self.assertRaises(MissingVariableError):
+            template.apply_vars({})
+
+    def test_template_apply_vars__no_matching_vars__exception(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        with self.assertRaises(MissingVariableError):
+            template.apply_vars({"not used": SOME_NON_TEMPLATE_TREE})
+
+    def test_template_apply_vars__matching_vars__template_replaced(self):
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+        template = Template(tree, conf=self.conf)
+
+        expected_result = deepcopy(tree)
+        expected_result.children[0] = SOME_NON_TEMPLATE_TREE
+        self.assertEqual(
+            expected_result,
+            template.apply_vars({SOME_TEMPLATE_NAME: SOME_NON_TEMPLATE_TREE}),
+        )
+
+
+class TestTreeTemplatesTemplateTranslator(unittest.TestCase):
+    parser = Lark(SOME_TEMPLATING_GRAMMAR)
+    conf = TemplateConf(parser.parse)
+
+    def test_translate__empty_translations__same_tree(self):
+        # no translations to match, so doesn't replace anything & can't error
+        translator = TemplateTranslator({})
+        tree = self.parser.parse(SOME_FOO_TEMPLATE)
+
+        expected_result = deepcopy(tree)
+        self.assertEqual(expected_result, translator.translate(tree))
+
+    def test_translate__one_translations__same_tree(self):
+        translations = {
+            self.conf(f"${SOME_TEMPLATE_NAME} bar"): self.conf(
+                f"--${SOME_TEMPLATE_NAME}-- bar"
+            )
+        }
+        translator = TemplateTranslator(translations)
+        tree = self.parser.parse(SOME_NON_TEMPLATED_STRING)
+
+        expected_result = deepcopy(tree)
+        expected_result.children.insert(0, Token("DASHES", "--"))
+        expected_result.children.insert(2, Token("DASHES", "--"))
+        self.assertEqual(expected_result, translator.translate(tree))
+
+
+class TestTreeTemplatesTemplateDefaultConf(unittest.TestCase):
+    def test_template_match__match_same_tree__empty_dictionary(self):
+        tree = Tree("foo", children=["hi"])
+        template = Template(tree)
+
+        self.assertEqual({}, template.match(tree))
+
+
+if __name__ == "__main__":
+    unittest.main()