Modified recursive dictionary parsing to correctly handle lists in dv…

…c plot templates (#134)

* Modified recursive dictionary parsing to correctly handle lists in dvc plot templates

* Rename dict_find_value -> find_value.

Extend tests.

* Update type hint

---------

Co-authored-by: daavoo <daviddelaiglesiacastro@gmail.com>

src/dvc_render/vega_templates.py

@@ -61,16 +61,21 @@ def list_replace_value(l: list, name: str, value: str) -> list:  # noqa: E741
     return x
 
 
-def dict_find_value(d: dict, value: str) -> bool:
-    for v in d.values():
-        if isinstance(v, dict):
-            if dict_find_value(v, value):
-                return True
-        if isinstance(v, str):
-            if v == value:
-                return True
-        if isinstance(v, list):
-            return any(dict_find_value(e, value) for e in v)
+def find_value(d: Union[dict, list, str], value: str) -> bool:
+    if isinstance(d, dict):
+        for v in d.values():
+            if isinstance(v, dict):
+                if find_value(v, value):
+                    return True
+            if isinstance(v, str):
+                if v == value:
+                    return True
+            if isinstance(v, list):
+                if any(find_value(e, value) for e in v):
+                    return True
+    elif isinstance(d, str):
+        if d == value:
+            return True
     return False
 
 
@@ -120,7 +125,7 @@ def reset(self):
 
     def has_anchor(self, name) -> bool:
         "Check if ANCHOR formatted with name is in content."
-        found = dict_find_value(self.content, self.anchor(name))
+        found = find_value(self.content, self.anchor(name))
         return found
 
     def fill_anchor(self, name, value) -> None:

tests/test_templates.py

@@ -10,8 +10,8 @@
     Template,
     TemplateContentDoesNotMatch,
     TemplateNotFoundError,
-    dict_find_value,
     dump_templates,
+    find_value,
     get_template,
 )
 
@@ -110,7 +110,8 @@ def test_escape_special_characters():
         ({"key": "value"}, "value"),
         ({"key": {"subkey": "value"}}, "value"),
         ({"key": [{"subkey": "value"}]}, "value"),
+        ({"key1": [{"subkey": "foo"}], "key2": {"subkey2": "value"}}, "value"),
     ],
 )
-def test_dict_find_value(content_dict, value_name):
-    assert dict_find_value(content_dict, value_name)
+def test_find_value(content_dict, value_name):
+    assert find_value(content_dict, value_name)

tests/test_vega.py

@@ -110,20 +110,35 @@ def test_bad_template_on_init():
         Template("name", "content")
 
 
-def test_bad_template_on_missing_data(tmp_dir):
-    template_content = {"data": {"values": "BAD_ANCHOR"}}
-    tmp_dir.gen("bar.json", json.dumps(template_content))
+@pytest.mark.parametrize(
+    "bad_content,good_content",
+    (
+        (
+            {"data": {"values": "BAD_ANCHOR"}},
+            {"data": {"values": Template.anchor("data")}},
+        ),
+        (
+            {"mark": {"type": "bar"}, "data": {"values": "BAD_ANCHOR"}},
+            {"mark": {"type": "bar"}, "data": {"values": Template.anchor("data")}},
+        ),
+        (
+            {"repeat": ["quintile"], "spec": {"data": {"values": "BAD_ANCHOR"}}},
+            {
+                "repeat": ["quintile"],
+                "spec": {"data": {"values": Template.anchor("data")}},
+            },
+        ),
+    ),
+)
+def test_bad_template_on_missing_data(tmp_dir, bad_content, good_content):
+    tmp_dir.gen("bar.json", json.dumps(bad_content))
     datapoints = [{"val": 2}, {"val": 3}]
     renderer = VegaRenderer(datapoints, "foo", template="bar.json")
 
     with pytest.raises(BadTemplateError):
         renderer.get_filled_template()
 
-    template_content = {
-        "mark": {"type": "bar"},
-        "data": {"values": Template.anchor("data")},
-    }
-    tmp_dir.gen("bar.json", json.dumps(template_content))
+    tmp_dir.gen("bar.json", json.dumps(good_content))
     renderer = VegaRenderer(datapoints, "foo", template="bar.json")
     assert renderer.get_filled_template()