html: Support passing custom_template as string. (#118)

daavoo · web-flow · commit 2aa47c86a4a3 · 2023-02-09T08:53:52.000+01:00
Will be used in DVCLive for treeverse/dvclive#460
diff --git a/src/dvc_render/html.py b/src/dvc_render/html.py
@@ -79,17 +79,18 @@ def embed(self) -> str:
 def render_html(
     renderers: List["Renderer"],
     output_file: "StrPath",
-    template_path: Optional["StrPath"] = None,
+    html_template: Optional["StrPath"] = None,
     refresh_seconds: Optional[int] = None,
 ) -> "StrPath":
-    "User renderers to fill an HTML template and write to path."
+    "Use `renderers` to fill an HTML template and write to `output_file`."
     output_path = Path(output_file)
     output_path.parent.mkdir(exist_ok=True)
 
-    page_html = None
-    if template_path:
-        with open(template_path, encoding="utf-8") as fobj:
-            page_html = fobj.read()
+    page_html: Optional[str] = None
+    if html_template and Path(html_template).is_file():
+        page_html = Path(html_template).read_text(encoding="utf8")
+    elif isinstance(html_template, str):
+        page_html = html_template
 
     document = HTML(page_html, refresh_seconds=refresh_seconds)
 
diff --git a/tests/test_html.py b/tests/test_html.py
@@ -1,9 +1,7 @@
-import pytest
-
-from dvc_render.html import HTML, PAGE_HTML, MissingPlaceholderError
-
 # pylint: disable=missing-function-docstring, R0801
+import pytest
 
+from dvc_render.html import HTML, PAGE_HTML, MissingPlaceholderError, render_html
 
 CUSTOM_PAGE_HTML = """<!DOCTYPE html>
 <html>
@@ -68,6 +66,23 @@ def test_html(template, page_elements, expected_page):
     assert result == expected_page
 
 
+def test_render_html_with_custom_template(mocker, tmp_dir):
+    output_file = tmp_dir / "output_file"
+
+    render_html(mocker.MagicMock(), output_file)
+    assert output_file.read_text() == PAGE_HTML.replace("{plot_divs}", "").replace(
+        "{scripts}", ""
+    ).replace("{refresh_tag}", "")
+
+    render_html(mocker.MagicMock(), output_file, CUSTOM_PAGE_HTML)
+    assert output_file.read_text() == CUSTOM_PAGE_HTML.format(plot_divs="")
+
+    custom_template = tmp_dir / "custom_template"
+    custom_template.write_text(CUSTOM_PAGE_HTML)
+    render_html(mocker.MagicMock(), output_file, custom_template)
+    assert output_file.read_text() == CUSTOM_PAGE_HTML.format(plot_divs="")
+
+
 def test_no_placeholder():
     template = "<head></head><body></body>"
 
