diff --git a/src/dvc_render/vega_templates.py b/src/dvc_render/vega_templates.py index a7af755..82fb527 100644 --- a/src/dvc_render/vega_templates.py +++ b/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: diff --git a/tests/test_templates.py b/tests/test_templates.py index 2fd88de..5c826e6 100644 --- a/tests/test_templates.py +++ b/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) diff --git a/tests/test_vega.py b/tests/test_vega.py index 1dbcb8b..f545665 100644 --- a/tests/test_vega.py +++ b/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()