add base64 support for markdown (#129)

* add base64 support for markdown

* return in-mem markdown if no report path

* fix mypy errors

src/dvc_render/image.py

@@ -55,7 +55,7 @@ def generate_markdown(self, report_path=None) -> str:
         for datapoint in self.datapoints:
             src = datapoint[self.SRC_FIELD]
             if src.startswith("data:image;base64"):
-                raise ValueError("`generate_markdown` doesn't support base64")
+                src = src.replace("data:image;base64", "data:image/png;base64")
             content.append(f"\n![{datapoint[self.TITLE_FIELD]}]({src})")
         if content:
             return "\n".join(content)

src/dvc_render/markdown.py

@@ -46,12 +46,14 @@ def embed(self) -> str:
 
 def render_markdown(
     renderers: List["Renderer"],
-    output_file: "StrPath",
+    output_file: Optional["StrPath"] = None,
     template_path: Optional["StrPath"] = None,
 ) -> "StrPath":
     "User renderers to fill an Markdown template and write to path."
-    output_path = Path(output_file)
-    output_path.parent.mkdir(exist_ok=True)
+    output_path = None
+    if output_file:
+        output_path = Path(output_file)
+        output_path.parent.mkdir(exist_ok=True)
 
     page = None
     if template_path:
@@ -63,6 +65,9 @@ def render_markdown(
     for renderer in renderers:
         document.with_element(renderer.generate_markdown(report_path=output_path))
 
-    output_path.write_text(document.embed(), encoding="utf8")
+    if output_file and output_path:
+        output_path.write_text(document.embed(), encoding="utf8")
 
-    return output_file
+        return output_file
+
+    return document.embed()

src/dvc_render/vega.py

@@ -1,3 +1,5 @@
+import base64
+import io
 import json
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -102,23 +104,35 @@ def generate_markdown(self, report_path=None) -> str:
 
         data = list_dict_to_dict_list(self.datapoints)
         if data:
-            report_folder = Path(report_path).parent
-            output_file = report_folder / self.name
-            output_file = output_file.with_suffix(".png")
-            output_file.parent.mkdir(exist_ok=True, parents=True)
+            if report_path:
+                report_folder = Path(report_path).parent
+                output_file = report_folder / self.name
+                output_file = output_file.with_suffix(".png")
+                output_file.parent.mkdir(exist_ok=True, parents=True)
+            else:
+                output_file = io.BytesIO()  # type: ignore
 
             x = self.properties.get("x")
             y = self.properties.get("y")
             data[x] = list(map(float, data[x]))
             data[y] = list(map(float, data[y]))
 
-            plt.title(self.properties.get("title", output_file.stem))
+            plt.title(self.properties.get("title", Path(self.name).stem))
             plt.xlabel(self.properties.get("x_label", x))
             plt.ylabel(self.properties.get("y_label", y))
             plt.plot(x, y, data=data)
             plt.tight_layout()
             plt.savefig(output_file)
             plt.close()
 
-            return f"\n![{self.name}]({output_file.relative_to(report_folder)})"
+            if report_path:
+                return f"\n![{self.name}]({output_file.relative_to(report_folder)})"
+
+            base64_str = base64.b64encode(
+                output_file.getvalue()  # type: ignore
+            ).decode()
+            src = f"data:image/png;base64,{base64_str}"
+
+            return f"\n![{self.name}]({src})"
+
         return ""

tests/test_image.py

@@ -63,8 +63,9 @@ def test_invalid_generate_markdown():
             "src": "data:image;base64,encoded_image",
         }
     ]
-    with pytest.raises(ValueError, match="`generate_markdown` doesn't support base64"):
-        ImageRenderer(datapoints, "file.jpg").generate_markdown()
+    md = ImageRenderer(datapoints, "file.jpg").generate_markdown()
+
+    assert "![workspace](data:image/png;base64,encoded_image)" in md
 
 
 @pytest.mark.parametrize(

tests/test_markdown.py

@@ -1,6 +1,11 @@
 import pytest
 
-from dvc_render.markdown import PAGE_MARKDOWN, Markdown, MissingPlaceholderError
+from dvc_render.markdown import (
+    PAGE_MARKDOWN,
+    Markdown,
+    MissingPlaceholderError,
+    render_markdown,
+)
 
 # pylint: disable=missing-function-docstring, R0801
 
@@ -26,7 +31,7 @@
         ),
     ],
 )
-def test_html(template, page_elements, expected_page):
+def test_markdown(template, page_elements, expected_page):
     page = Markdown(template)
     page.elements = page_elements
 
@@ -40,3 +45,12 @@ def test_no_placeholder():
 
     with pytest.raises(MissingPlaceholderError):
         Markdown(template)
+
+
+def test_render_markdown_to_file(tmp_dir):
+    output_file = tmp_dir / "report"
+    assert output_file == render_markdown([], output_file)
+
+
+def test_render_markdown_no_file():
+    assert "# DVC Report" in render_markdown([])

tests/test_vega.py

@@ -115,7 +115,9 @@ def test_raise_on_wrong_field():
 
 
 @pytest.mark.parametrize("name", ["foo", "foo/bar", "foo/bar.tsv"])
-def test_generate_markdown(tmp_dir, mocker, name):
+@pytest.mark.parametrize("to_file", [True, False])
+def test_generate_markdown(tmp_dir, mocker, name, to_file):
+    # pylint: disable-msg=too-many-locals
     import matplotlib.pyplot
 
     plot = mocker.spy(matplotlib.pyplot, "plot")
@@ -131,10 +133,18 @@ def test_generate_markdown(tmp_dir, mocker, name):
     ]
     renderer = VegaRenderer(datapoints, name, **props)
 
-    (tmp_dir / "output").mkdir()
-    renderer.generate_markdown(tmp_dir / "output" / "report.md")
+    if to_file:
+        report_folder = tmp_dir / "output"
+        report_folder.mkdir()
+        md = renderer.generate_markdown(tmp_dir / "output" / "report.md")
+        output_file = (tmp_dir / "output" / renderer.name).with_suffix(".png")
+        assert output_file.exists()
+        savefig.assert_called_with(output_file)
+        assert f"![{name}]({output_file.relative_to(report_folder)})" in md
+    else:
+        md = renderer.generate_markdown()
+        assert f"![{name}](data:image/png;base64," in md
 
-    assert (tmp_dir / "output" / renderer.name).with_suffix(".png").exists()
     plot.assert_called_with(
         "first_val",
         "second_val",
@@ -147,7 +157,6 @@ def test_generate_markdown(tmp_dir, mocker, name):
     title.assert_called_with("FOO")
     xlabel.assert_called_with("first_val")
     ylabel.assert_called_with("second_val")
-    savefig.assert_called_with((tmp_dir / "output" / name).with_suffix(".png"))
 
 
 def test_unsupported_template():