From 49b8f8a81c4e06b8f675197b8dd57e2a773cf283 Mon Sep 17 00:00:00 2001
From: David de la Iglesia Castro <daviddelaiglesiacastro@gmail.com>
Date: Thu, 16 Mar 2023 20:28:22 +0100
Subject: [PATCH] vega: Add `as_string` argument. Default True. (#125)

`0.3.0` release broke `dvc plots --json` for older DVC installations because `dvc-render` dep didn't had an upper bound.
---
 src/dvc_render/vega.py | 14 ++++++++++----
 tests/test_vega.py     | 12 ++++++++----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/dvc_render/vega.py b/src/dvc_render/vega.py
index 7d1f4cd..6848eeb 100644
--- a/src/dvc_render/vega.py
+++ b/src/dvc_render/vega.py
@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from warnings import warn
 
 from .base import Renderer
@@ -38,8 +38,11 @@ def __init__(self, datapoints: List, name: str, **properties):
         )
 
     def get_filled_template(
-        self, skip_anchors: Optional[List[str]] = None, strict: bool = True
-    ) -> Dict[str, Any]:
+        self,
+        skip_anchors: Optional[List[str]] = None,
+        strict: bool = True,
+        as_string: bool = True,
+    ) -> Union[str, Dict[str, Any]]:
         """Returns a functional vega specification"""
         self.template.reset()
         if not self.datapoints:
@@ -80,10 +83,13 @@ def get_filled_template(
                 value = self.template.escape_special_characters(value)
             self.template.fill_anchor(name, value)
 
+        if as_string:
+            return json.dumps(self.template.content)
+
         return self.template.content
 
     def partial_html(self, **kwargs) -> str:
-        return json.dumps(self.get_filled_template())
+        return self.get_filled_template()  # type: ignore
 
     def generate_markdown(self, report_path=None) -> str:
         if not isinstance(self.template, LinearTemplate):
diff --git a/tests/test_vega.py b/tests/test_vega.py
index dbc823b..0aced9f 100644
--- a/tests/test_vega.py
+++ b/tests/test_vega.py
@@ -40,7 +40,7 @@ def test_default_template_mark():
         {"first_val": 200, "second_val": 300, "val": 3},
     ]
 
-    plot_content = VegaRenderer(datapoints, "foo").get_filled_template()
+    plot_content = VegaRenderer(datapoints, "foo").get_filled_template(as_string=False)
 
     assert plot_content["layer"][0]["mark"] == "line"
 
@@ -57,7 +57,9 @@ def test_choose_axes():
         {"first_val": 200, "second_val": 300, "val": 3},
     ]
 
-    plot_content = VegaRenderer(datapoints, "foo", **props).get_filled_template()
+    plot_content = VegaRenderer(datapoints, "foo", **props).get_filled_template(
+        as_string=False
+    )
 
     assert plot_content["data"]["values"] == [
         {
@@ -82,7 +84,9 @@ def test_confusion():
     ]
     props = {"template": "confusion", "x": "predicted", "y": "actual"}
 
-    plot_content = VegaRenderer(datapoints, "foo", **props).get_filled_template()
+    plot_content = VegaRenderer(datapoints, "foo", **props).get_filled_template(
+        as_string=False
+    )
 
     assert plot_content["data"]["values"] == [
         {"predicted": "B", "actual": "A"},
@@ -170,7 +174,7 @@ def test_escape_special_characters():
     ]
     props = {"template": "simple", "x": "foo.bar[0]", "y": "foo.bar[1]"}
     renderer = VegaRenderer(datapoints, "foo", **props)
-    filled = renderer.get_filled_template()
+    filled = renderer.get_filled_template(as_string=False)
     # data is not escaped
     assert filled["data"]["values"][0] == datapoints[0]
     # field and title yes