diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py index f178ab51e..2b4ed1aea 100644 --- a/dvc/repo/plots/__init__.py +++ b/dvc/repo/plots/__init__.py @@ -187,7 +187,7 @@ class Plots: props = props or {} template = props.get("template") if template: - self.templates.load(template) + self.templates.get_template(template) (out,) = self.repo.find_outs_by_path(path) if not out.plot and unset is not None: diff --git a/dvc/repo/plots/template.py b/dvc/repo/plots/template.py index 9c8dffde2..69b94e43f 100644 --- a/dvc/repo/plots/template.py +++ b/dvc/repo/plots/template.py @@ -1,23 +1,15 @@ import json import os -from typing import TYPE_CHECKING, Iterable, Optional - -try: - import importlib_resources -except ImportError: - import importlib.resources as importlib_resources # type: ignore[no-redef] +from typing import Any, Dict, Optional from funcy import cached_property from dvc.exceptions import DvcException -if TYPE_CHECKING: - from dvc.types import StrPath - class TemplateNotFoundError(DvcException): - def __init__(self, name): - super().__init__(f"Template '{name}' not found.") + def __init__(self, path): + super().__init__(f"Template '{path}' not found.") class NoFieldInDataError(DvcException): @@ -33,9 +25,25 @@ class Template: EXTENSION = ".json" ANCHOR = "" - def __init__(self, content, name): - self.content = content - self.name = name + DEFAULT_CONTENT: Optional[Dict[str, Any]] = None + DEFAULT_NAME: Optional[str] = None + + def __init__(self, content=None, name=None): + if content: + self.content = content + else: + self.content = ( + json.dumps( + self.DEFAULT_CONTENT, + indent=self.INDENT, + separators=self.SEPARATORS, + ) + + "\n" + ) + + self.name = name or self.DEFAULT_NAME + assert self.content and self.name + self.filename = self.name + self.EXTENSION @classmethod def anchor(cls, name): @@ -61,86 +69,466 @@ class Template: raise NoFieldInDataError(field) +class SimpleLinearTemplate(Template): + DEFAULT_NAME = "simple" + + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "width": 300, + "height": 300, + "mark": {"type": "line"}, + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + "title": Template.anchor("y_label"), + "scale": {"zero": False}, + }, + "color": {"field": "rev", "type": "nominal"}, + }, + } + + +class ConfusionTemplate(Template): + DEFAULT_NAME = "confusion" + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "facet": {"field": "rev", "type": "nominal"}, + "spec": { + "transform": [ + { + "aggregate": [{"op": "count", "as": "xy_count"}], + "groupby": [Template.anchor("y"), Template.anchor("x")], + }, + { + "impute": "xy_count", + "groupby": ["rev", Template.anchor("y")], + "key": Template.anchor("x"), + "value": 0, + }, + { + "impute": "xy_count", + "groupby": ["rev", Template.anchor("x")], + "key": Template.anchor("y"), + "value": 0, + }, + { + "joinaggregate": [ + {"op": "max", "field": "xy_count", "as": "max_count"} + ], + "groupby": [], + }, + { + "calculate": "datum.xy_count / datum.max_count", + "as": "percent_of_max", + }, + ], + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "nominal", + "sort": "ascending", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "nominal", + "sort": "ascending", + "title": Template.anchor("y_label"), + }, + }, + "layer": [ + { + "mark": "rect", + "width": 300, + "height": 300, + "encoding": { + "color": { + "field": "xy_count", + "type": "quantitative", + "title": "", + "scale": {"domainMin": 0, "nice": True}, + } + }, + }, + { + "mark": "text", + "encoding": { + "text": {"field": "xy_count", "type": "quantitative"}, + "color": { + "condition": { + "test": "datum.percent_of_max > 0.5", + "value": "white", + }, + "value": "black", + }, + }, + }, + ], + }, + } + + +class NormalizedConfusionTemplate(Template): + DEFAULT_NAME = "confusion_normalized" + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "facet": {"field": "rev", "type": "nominal"}, + "spec": { + "transform": [ + { + "aggregate": [{"op": "count", "as": "xy_count"}], + "groupby": [Template.anchor("y"), Template.anchor("x")], + }, + { + "impute": "xy_count", + "groupby": ["rev", Template.anchor("y")], + "key": Template.anchor("x"), + "value": 0, + }, + { + "impute": "xy_count", + "groupby": ["rev", Template.anchor("x")], + "key": Template.anchor("y"), + "value": 0, + }, + { + "joinaggregate": [ + {"op": "sum", "field": "xy_count", "as": "sum_y"} + ], + "groupby": [Template.anchor("y")], + }, + { + "calculate": "datum.xy_count / datum.sum_y", + "as": "percent_of_y", + }, + ], + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "nominal", + "sort": "ascending", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "nominal", + "sort": "ascending", + "title": Template.anchor("y_label"), + }, + }, + "layer": [ + { + "mark": "rect", + "width": 300, + "height": 300, + "encoding": { + "color": { + "field": "percent_of_y", + "type": "quantitative", + "title": "", + "scale": {"domain": [0, 1]}, + } + }, + }, + { + "mark": "text", + "encoding": { + "text": { + "field": "percent_of_y", + "type": "quantitative", + "format": ".2f", + }, + "color": { + "condition": { + "test": "datum.percent_of_y > 0.5", + "value": "white", + }, + "value": "black", + }, + }, + }, + ], + }, + } + + +class ScatterTemplate(Template): + DEFAULT_NAME = "scatter" + + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "width": 300, + "height": 300, + "layer": [ + { + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + "title": Template.anchor("y_label"), + "scale": {"zero": False}, + }, + "color": {"field": "rev", "type": "nominal"}, + }, + "layer": [ + {"mark": "point"}, + { + "selection": { + "label": { + "type": "single", + "nearest": True, + "on": "mouseover", + "encodings": ["x"], + "empty": "none", + "clear": "mouseout", + } + }, + "mark": "point", + "encoding": { + "opacity": { + "condition": { + "selection": "label", + "value": 1, + }, + "value": 0, + } + }, + }, + ], + }, + { + "transform": [{"filter": {"selection": "label"}}], + "layer": [ + { + "encoding": { + "text": { + "type": "quantitative", + "field": Template.anchor("y"), + }, + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + }, + }, + "layer": [ + { + "mark": { + "type": "text", + "align": "left", + "dx": 5, + "dy": -5, + }, + "encoding": { + "color": { + "type": "nominal", + "field": "rev", + } + }, + } + ], + } + ], + }, + ], + } + + +class SmoothLinearTemplate(Template): + DEFAULT_NAME = "smooth" + + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "mark": {"type": "line"}, + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + "title": Template.anchor("y_label"), + "scale": {"zero": False}, + }, + "color": {"field": "rev", "type": "nominal"}, + }, + "transform": [ + { + "loess": Template.anchor("y"), + "on": Template.anchor("x"), + "groupby": ["rev"], + "bandwidth": 0.3, + } + ], + } + + +class LinearTemplate(Template): + DEFAULT_NAME = "linear" + + DEFAULT_CONTENT = { + "$schema": "https://vega.github.io/schema/vega-lite/v4.json", + "data": {"values": Template.anchor("data")}, + "title": Template.anchor("title"), + "width": 300, + "height": 300, + "layer": [ + { + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + "title": Template.anchor("x_label"), + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + "title": Template.anchor("y_label"), + "scale": {"zero": False}, + }, + "color": {"field": "rev", "type": "nominal"}, + }, + "layer": [ + {"mark": "line"}, + { + "selection": { + "label": { + "type": "single", + "nearest": True, + "on": "mouseover", + "encodings": ["x"], + "empty": "none", + "clear": "mouseout", + } + }, + "mark": "point", + "encoding": { + "opacity": { + "condition": { + "selection": "label", + "value": 1, + }, + "value": 0, + } + }, + }, + ], + }, + { + "transform": [{"filter": {"selection": "label"}}], + "layer": [ + { + "mark": {"type": "rule", "color": "gray"}, + "encoding": { + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + } + }, + }, + { + "encoding": { + "text": { + "type": "quantitative", + "field": Template.anchor("y"), + }, + "x": { + "field": Template.anchor("x"), + "type": "quantitative", + }, + "y": { + "field": Template.anchor("y"), + "type": "quantitative", + }, + }, + "layer": [ + { + "mark": { + "type": "text", + "align": "left", + "dx": 5, + "dy": -5, + }, + "encoding": { + "color": { + "type": "nominal", + "field": "rev", + } + }, + } + ], + }, + ], + }, + ], + } + + class PlotTemplates: TEMPLATES_DIR = "plots" - PKG_TEMPLATES_DIR = "templates" + TEMPLATES = [ + SimpleLinearTemplate, + LinearTemplate, + ConfusionTemplate, + NormalizedConfusionTemplate, + ScatterTemplate, + SmoothLinearTemplate, + ] @cached_property def templates_dir(self): return os.path.join(self.dvc_dir, self.TEMPLATES_DIR) - @staticmethod - def _find(templates, template_name): - for template in templates: - if template.endswith(template_name) or template.endswith( - template_name + ".json" - ): - return template - return None + def get_template(self, template_name: str): + template_path = os.path.abspath(template_name) + if os.path.exists(template_path): + return os.path.abspath(template_path) - def _find_in_project(self, name: str) -> Optional["StrPath"]: - full_path = os.path.abspath(name) - if os.path.exists(full_path): - return full_path + if self.dvc_dir and os.path.exists(self.dvc_dir): + template_path = os.path.join(self.templates_dir, template_name) + if os.path.exists(template_path): + return template_path - if os.path.exists(self.templates_dir): - templates = [ + all_templates = [ os.path.join(root, file) for root, _, files in os.walk(self.templates_dir) for file in files ] - found = self._find(templates, name) - if found: - return os.path.join(self.templates_dir, found) - return None - - @staticmethod - def _get_templates() -> Iterable[str]: - if ( - importlib_resources.files(__package__) - .joinpath(PlotTemplates.PKG_TEMPLATES_DIR) - .is_dir() - ): - entries = ( - importlib_resources.files(__package__) - .joinpath(PlotTemplates.PKG_TEMPLATES_DIR) - .iterdir() - ) - return [entry.name for entry in entries] - return [] - - @staticmethod - def _load_from_pkg(name): - templates = PlotTemplates._get_templates() - found = PlotTemplates._find(templates, name) - if found: - return ( - ( - importlib_resources.files(__package__) - / PlotTemplates.PKG_TEMPLATES_DIR - / found - ) - .read_bytes() - .decode("utf-8") - ) - return None - - def load(self, name: str = None) -> Template: - - if name is not None: - template_path = self._find_in_project(name) - if template_path: - with open(template_path, "r", encoding="utf-8") as fd: - content = fd.read() - return Template(content, name) - else: - name = "linear" - - content = self._load_from_pkg(name) - if content: - return Template(content, name) - - raise TemplateNotFoundError(name) + matches = [ + template + for template in all_templates + if os.path.splitext(template)[0] == template_path + ] + if matches: + assert len(matches) == 1 + return matches[0] + raise TemplateNotFoundError(template_name) def __init__(self, dvc_dir): self.dvc_dir = dvc_dir @@ -149,18 +537,27 @@ class PlotTemplates: from dvc.utils.fs import makedirs makedirs(self.templates_dir, exist_ok=True) + for t in self.TEMPLATES: + self._dump(t()) - templates = self._get_templates() - for template in templates: - content = ( - importlib_resources.files(__package__) - .joinpath(PlotTemplates.PKG_TEMPLATES_DIR) - .joinpath(template) - .read_text() - ) - with open( - os.path.join(self.templates_dir, template), - "w", - encoding="utf-8", - ) as fd: - fd.write(content) + def _dump(self, template): + path = os.path.join(self.templates_dir, template.filename) + with open(path, "w", encoding="utf-8") as fd: + fd.write(template.content) + + def load(self, template_name=None): + if not template_name: + template_name = "linear" + + try: + path = self.get_template(template_name) + + with open(path, "r", encoding="utf-8") as fd: + content = fd.read() + + return Template(content, name=template_name) + except TemplateNotFoundError: + for template in self.TEMPLATES: + if template.DEFAULT_NAME == template_name: + return template() + raise diff --git a/tests/unit/render/test_vega.py b/tests/unit/render/test_vega.py index ca9dad68b..7ae179c5c 100644 --- a/tests/unit/render/test_vega.py +++ b/tests/unit/render/test_vega.py @@ -471,6 +471,6 @@ def test_should_resolve_template(tmp_dir, dvc, template_path, target_name): with open(template_path, "w", encoding="utf-8") as fd: fd.write("template_content") - assert dvc.plots.templates._find_in_project( - target_name - ) == os.path.abspath(template_path) + assert dvc.plots.templates.get_template(target_name) == os.path.abspath( + template_path + )