treeverse · pared · Apr 20, 2021 · Mar 25, 2021 · Apr 12, 2021 · jorgeorpinel
diff --git a/dvc/command/live.py b/dvc/command/live.py
@@ -4,7 +4,6 @@
 
 from dvc.command import completion
 from dvc.command.base import CmdBase, fix_subparsers
-from dvc.utils.html import write
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +15,7 @@ def _run(self, target, revs=None):
         metrics, plots = self.repo.live.show(target=target, revs=revs)
 
         html_path = self.args.target + ".html"
-        write(html_path, plots, metrics)
+        self.repo.plots.write_html(html_path, plots, metrics)
 
         logger.info(f"\nfile://{os.path.abspath(html_path)}")
 

diff --git a/dvc/command/plots.py b/dvc/command/plots.py
@@ -5,7 +5,6 @@
 from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
 from dvc.exceptions import DvcException
 from dvc.utils import format_link
-from dvc.utils.html import write
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +43,10 @@ def run(self):
 
             rel: str = self.args.out or "plots.html"
             path = (Path.cwd() / rel).resolve()
-            write(path, plots)
+            self.repo.plots.write_html(
+                path, plots=plots, html_template_path=self.args.html_template
+            )
+
         except DvcException:
             logger.exception("")
             return 1
@@ -243,3 +245,9 @@ def _add_output_arguments(parser):
         default=False,
         help="Open plot file directly in the browser.",
     )
+    parser.add_argument(
+        "--html-template",
+        default=None,
+        help="Custom HTML template for VEGA visualization.",
+        metavar="<path>",
+    )
diff --git a/dvc/config_schema.py b/dvc/config_schema.py
@@ -223,4 +223,5 @@ class RelPath(str):
         # enabled by default. It's of no use, kept for backward compatibility.
         Optional("parametrization", default=True): Bool
     },
+    "plots": {"html_template": str},
 }
diff --git a/dvc/repo/live.py b/dvc/repo/live.py
@@ -14,14 +14,13 @@
 
 
 def create_summary(out):
-    from dvc.utils.html import write
-
     assert out.live and out.live["html"]
 
     metrics, plots = out.repo.live.show(str(out.path_info))
 
     html_path = out.path_info.with_suffix(".html")
-    write(html_path, plots, metrics)
+
+    out.repo.plots.write_html(html_path, plots, metrics)
     logger.info(f"\nfile://{os.path.abspath(html_path)}")
 
 

diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py
@@ -1,5 +1,6 @@
 import logging
-from typing import TYPE_CHECKING, Dict, List
+import os
+from typing import TYPE_CHECKING, Dict, List, Optional
 
 from funcy import cached_property, first, project
 
@@ -9,6 +10,7 @@
     NoMetricsFoundError,
     NoMetricsParsedError,
 )
+from dvc.types import StrPath
 from dvc.utils import relpath
 
 if TYPE_CHECKING:
@@ -219,6 +221,26 @@ def templates(self):
 
         return PlotTemplates(self.repo.dvc_dir)
 
+    def write_html(
+        self,
+        path: StrPath,
+        plots: Dict[str, Dict],
+        metrics: Optional[Dict[str, Dict]] = None,
+        html_template_path: Optional[StrPath] = None,
+    ):
+        if not html_template_path:
+            html_template_path = self.repo.config.get("plots", {}).get(
+                "html_template", None
+            )
+            if html_template_path and not os.path.isabs(html_template_path):
+                html_template_path = os.path.join(
+                    self.repo.dvc_dir, html_template_path
+                )
+
+        from dvc.utils.html import write
+
+        write(path, plots, metrics, html_template_path)
+
 
 def _is_plot(out: "BaseOutput") -> bool:
     return bool(out.plot) or bool(out.live)

diff --git a/dvc/utils/html.py b/dvc/utils/html.py
@@ -1,5 +1,8 @@
 from typing import Dict, List, Optional
 
+from dvc.exceptions import DvcException
+from dvc.types import StrPath
+
 PAGE_HTML = """<!DOCTYPE html>
 <html>
 <head>
@@ -9,7 +12,7 @@
     <script src="https://cdn.jsdelivr.net/npm/vega-embed@6.5.1"></script>
 </head>
 <body>
-    {divs}
+    {plot_divs}
 </body>
 </html>"""
 
@@ -20,9 +23,22 @@
 </script>"""
 
 
+class MissingPlaceholderError(DvcException):
+    def __init__(self, placeholder):
+        super().__init__(f"HTML template has to contain '{placeholder}'.")
+
+
 class HTML:
-    def __init__(self):
-        self.elements = []
+    PLACEHOLDER = "plot_divs"
+    PLACEHOLDER_FORMAT_STR = f"{{{PLACEHOLDER}}}"
+
+    def __init__(self, template: str = None):
+        template = template or PAGE_HTML
+        if self.PLACEHOLDER_FORMAT_STR not in template:
+            raise MissingPlaceholderError(self.PLACEHOLDER_FORMAT_STR)
+
+        self.template = template
+        self.elements: List[str] = []
 
     def with_metrics(self, metrics: Dict[str, Dict]) -> "HTML":
         import tabulate
@@ -54,13 +70,22 @@ def with_element(self, html: str) -> "HTML":
         return self
 
     def embed(self) -> str:
-        return PAGE_HTML.format(divs="\n".join(self.elements))
+        kwargs = {self.PLACEHOLDER: "\n".join(self.elements)}
+        return self.template.format(**kwargs)
 
 
 def write(
-    path, plots: Dict[str, Dict], metrics: Optional[Dict[str, Dict]] = None
+    path: StrPath,
+    plots: Dict[str, Dict],
+    metrics: Optional[Dict[str, Dict]] = None,
+    template_path: Optional[StrPath] = None,
 ):
-    document = HTML()
+    page_html = None
+    if template_path:
+        with open(template_path, "r") as fobj:
+            page_html = fobj.read()
+
+    document = HTML(page_html)
     if metrics:
         document.with_metrics(metrics)
         document.with_element("<br>")

diff --git a/tests/func/utils/test_html.py b/tests/func/utils/test_html.py
@@ -0,0 +1,43 @@
+import pytest
+
+from dvc.utils.html import HTML, PAGE_HTML, MissingPlaceholderError
+
+CUSTOM_PAGE_HTML = """<!DOCTYPE html>
+<html>
+<head>
+    <title>TITLE</title>
+    <script type="text/javascript" src="vega"></script>
+    <script type="text/javascript" src="vega-lite"></script>
+    <script type="text/javascript" src="vega-embed"></script>
+</head>
+<body>
+    {plot_divs}
+</body>
+</html>"""
+
+
+@pytest.mark.parametrize(
+    "template,page_elements,expected_page",
+    [
+        (None, ["content"], PAGE_HTML.format(plot_divs="content")),
+        (
+            CUSTOM_PAGE_HTML,
+            ["content"],
+            CUSTOM_PAGE_HTML.format(plot_divs="content"),
+        ),
+    ],
+)
+def test_html(tmp_dir, template, page_elements, expected_page):
+    page = HTML(template)
+    page.elements = page_elements
+
+    result = page.embed()
+
+    assert result == expected_page
+
+
+def test_no_placeholder(tmp_dir):
+    template = "<head></head><body></body>"
+
+    with pytest.raises(MissingPlaceholderError):
+        HTML(template)