From 3b4d76df4285462ad896562f0b4b74aef307214f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Thu, 10 Feb 2022 01:22:16 +0100 Subject: [PATCH] plots: json: add datapoints to vega plots --- dvc/commands/plots.py | 14 ++++++--- dvc/render/data.py | 7 ++--- dvc/render/utils.py | 11 ++----- dvc/render/vega.py | 54 +++++++++++++++++++++++--------- tests/unit/command/test_plots.py | 3 +- tests/unit/render/test_data.py | 10 +++--- tests/unit/render/test_vega.py | 18 +++++++++++ 7 files changed, 82 insertions(+), 35 deletions(-) diff --git a/dvc/commands/plots.py b/dvc/commands/plots.py index 13aa9f5c3..62ee6ebc3 100644 --- a/dvc/commands/plots.py +++ b/dvc/commands/plots.py @@ -14,12 +14,14 @@ from dvc.utils import format_link logger = logging.getLogger(__name__) -def _show_json(renderers, path: None): +def _show_json(renderers, path: None, split=False): if any(r.needs_output_path for r in renderers) and not path: raise DvcException("Output path ('-o') is required!") result = { - renderer.filename: json.loads(renderer.as_json(path=path)) + renderer.filename: json.loads( + renderer.as_json(path=path, fill_data=not split) + ) for renderer in renderers } if result: @@ -80,10 +82,11 @@ class CmdPlots(CmdBase): filter(lambda r: isinstance(r, VegaRenderer), renderers) ) if renderer: - ui.write_json(renderer.asdict()) + content = renderer.asdict() + ui.write_json(content) return 0 if self.args.json: - _show_json(renderers, self.args.out) + _show_json(renderers, self.args.out, self.args.split) return 0 rel: str = self.args.out or "dvc_plots" @@ -354,6 +357,9 @@ def _add_ui_arguments(parser): default=False, help=argparse.SUPPRESS, ) + parser.add_argument( + "--split", action="store_true", default=False, help=argparse.SUPPRESS + ) parser.add_argument( "--open", action="store_true", diff --git a/dvc/render/data.py b/dvc/render/data.py index 5f45720ff..9ff3d09f6 100644 --- a/dvc/render/data.py +++ b/dvc/render/data.py @@ -180,15 +180,14 @@ class Converter: def to_datapoints(data: Dict, props: Dict): converter = Converter(props) - datapoints = [] + datapoints: Dict[str, List[Dict]] = {} for revision, rev_data in data.items(): + datapoints[revision] = [] for _, file_data in rev_data.get("data", {}).items(): if "data" in file_data: processed, final_props = converter.convert( file_data.get("data") ) - Converter.update(processed, {REVISION_FIELD: revision}) - - datapoints.extend(processed) + datapoints[revision].extend(processed) return datapoints, final_props diff --git a/dvc/render/utils.py b/dvc/render/utils.py index a5bdc0ac6..e27b20c1c 100644 --- a/dvc/render/utils.py +++ b/dvc/render/utils.py @@ -4,13 +4,9 @@ from typing import Dict, List import dpath.util -def get_files(data: Dict) -> List: - files = set() - for rev in data.keys(): - for file in data[rev].get("data", {}).keys(): - files.add(file) - sorted_files = sorted(files) - return sorted_files +def get_files(plots_data: Dict) -> List: + values = dpath.util.values(plots_data, ["*", "*"]) + return sorted({key for submap in values for key in submap.keys()}) def group_by_filename(plots_data: Dict) -> List[Dict]: @@ -35,7 +31,6 @@ def match_renderers(plots_data, templates): renderers = [] for group in group_by_filename(plots_data): - plot_properties = squash_plots_properties(group) template = templates.load(plot_properties.get("template", None)) diff --git a/dvc/render/vega.py b/dvc/render/vega.py index a7ad894bc..f44484481 100644 --- a/dvc/render/vega.py +++ b/dvc/render/vega.py @@ -1,14 +1,25 @@ import json import os from copy import deepcopy -from typing import Dict, Optional +from typing import Dict, List, Optional -from dvc.render.base import BadTemplateError, Renderer -from dvc.render.data import to_datapoints +from funcy import cached_property + +from dvc.render.base import REVISION_FIELD, BadTemplateError, Renderer +from dvc.render.data import Converter, to_datapoints from dvc.render.utils import get_files from dvc.repo.plots.template import Template +def _flatten_datapoints(rev_datapoints: Dict[str, List[Dict]]) -> List[Dict]: + flat = [] + for revision, datapoints in rev_datapoints.items(): + datapoints_cp = deepcopy(datapoints) + Converter.update(datapoints_cp, {REVISION_FIELD: revision}) + flat.extend(datapoints_cp) + return flat + + class VegaRenderer(Renderer): TYPE = "vega" @@ -37,7 +48,8 @@ class VegaRenderer(Renderer): def _revisions(self): return list(self.data.keys()) - def _fill_template(self, template, datapoints, props=None): + # TODO Shouldn't it be a part of plots? + def _fill_template(self, template, datapoints, props=None, fill_data=True): props = props or {} content = deepcopy(template.content) @@ -52,7 +64,8 @@ class VegaRenderer(Renderer): if props.get("y"): template.check_field_exists(datapoints, props.get("y")) - content = template.fill_anchor(content, "data", datapoints) + if fill_data: + content = template.fill_anchor(content, "data", datapoints) props.setdefault("title", "") props.setdefault("x_label", props.get("x")) @@ -66,27 +79,38 @@ class VegaRenderer(Renderer): return content - def get_filled_template(self): - props = self.properties - datapoints, final_props = to_datapoints(self.data, props) + @cached_property + def _converted_data(self): + return to_datapoints(self.data, self.properties) + + def _get_filled_template(self, fill_data=True): + datapoints, final_props = self._converted_data + flat_datapoints = _flatten_datapoints(datapoints) if datapoints: filled_template = self._fill_template( - self.template, datapoints, final_props + self.template, + flat_datapoints, + final_props, + fill_data=fill_data, ) return filled_template - return None + return "" - def asdict(self): - filled_template = self.get_filled_template() + # TODO rename + def asdict(self, fill_data=True): + filled_template = self._get_filled_template(fill_data) if filled_template: return json.loads(filled_template) return {} def as_json(self, **kwargs) -> Optional[str]: - content = self.asdict() + fill_data = kwargs.get("fill_data", True) + + content = self.asdict(fill_data) + datapoints, _ = self._converted_data return json.dumps( [ @@ -94,12 +118,14 @@ class VegaRenderer(Renderer): self.TYPE_KEY: self.TYPE, self.REVISIONS_KEY: self._revisions(), "content": content, + "datapoints": datapoints, } ], ) def partial_html(self, **kwargs): - return self.get_filled_template() or "" + filled_template = self._get_filled_template() + return filled_template @staticmethod def matches(data): diff --git a/tests/unit/command/test_plots.py b/tests/unit/command/test_plots.py index 6b855b184..124d3bb84 100644 --- a/tests/unit/command/test_plots.py +++ b/tests/unit/command/test_plots.py @@ -276,6 +276,7 @@ def test_plots_diff_json(dvc, mocker, capsys): "HEAD~10", "HEAD~1", "--json", + "--split", "--targets", "plot.csv", "-o", @@ -295,7 +296,7 @@ def test_plots_diff_json(dvc, mocker, capsys): assert cmd.run() == 0 - show_json_mock.assert_called_once_with(renderers, "out") + show_json_mock.assert_called_once_with(renderers, "out", True) render_mock.assert_not_called() diff --git a/tests/unit/render/test_data.py b/tests/unit/render/test_data.py index fc36c61db..3559b2962 100644 --- a/tests/unit/render/test_data.py +++ b/tests/unit/render/test_data.py @@ -150,10 +150,12 @@ def test_to_datapoints(): datapoints, resolved_properties = to_datapoints(input_data, props) - assert datapoints == [ - {"v": 1, "v2": 0.1, "v3": 0.01, "rev": "revision"}, - {"v": 2, "v2": 0.2, "v3": 0.02, "rev": "revision"}, - ] + assert datapoints == { + "revision": [ + {"v": 1, "v2": 0.1, "v3": 0.01}, + {"v": 2, "v2": 0.2, "v3": 0.02}, + ] + } assert resolved_properties == { "fields": {"v", "v2", "v3"}, "x": "v2", diff --git a/tests/unit/render/test_vega.py b/tests/unit/render/test_vega.py index 324f20792..f9d3e996b 100644 --- a/tests/unit/render/test_vega.py +++ b/tests/unit/render/test_vega.py @@ -341,6 +341,7 @@ def test_should_resolve_template(tmp_dir, dvc, template_path, target_name): ) +# TODO remove content checks def test_as_json(tmp_dir, scm, dvc): metric = [{"a": 1, "b": 2, "c": 3}, {"a": 2, "b": 3, "c": 4}] data = {"workspace": {"data": {"file.json": {"data": metric}}}} @@ -355,3 +356,20 @@ def test_as_json(tmp_dir, scm, dvc): assert plot_as_json["type"] == "vega" assert plot_as_json["revisions"] == ["workspace"] assert plot_as_json["content"] == plot_content + + +def test_as_json_split(tmp_dir, scm, dvc): + metric = [{"a": 1, "b": 2, "c": 3}, {"a": 2, "b": 3, "c": 4}] + data = {"workspace": {"data": {"file.json": {"data": metric}}}} + props = {"fields": {"b"}} + + renderer = VegaRenderer( + data, template=dvc.plots.templates.load(), properties=props + ) + plot_content = renderer.asdict(fill_data=False) + plot_as_json = first(json.loads(renderer.as_json(fill_data=False))) + + assert plot_as_json["type"] == "vega" + assert plot_as_json["revisions"] == ["workspace"] + assert plot_as_json["content"] == plot_content + assert plot_as_json["datapoints"] == renderer._converted_data[0]