Skip to content

Commit

Permalink
plots: make generator out of collect (iterative#6793)
Browse files Browse the repository at this point in the history
  • Loading branch information
pared authored Oct 15, 2021
1 parent 7f5c53c commit db75878
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 29 deletions.
51 changes: 27 additions & 24 deletions dvc/repo/plots/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import logging
import os
from collections import OrderedDict
from typing import TYPE_CHECKING, Callable, Dict, List, Optional
from typing import TYPE_CHECKING, Callable, Dict, Generator, List, Optional

from funcy import cached_property, first, project

Expand Down Expand Up @@ -47,10 +47,10 @@ def collect(
recursive: bool = False,
onerror: Optional[Callable] = None,
props: Optional[Dict] = None,
) -> Dict[str, Dict]:
) -> Generator[Dict, None, None]:
"""Collects all props and data for plots.
Returns a structure like:
Generator yielding a structure like:
{rev: {plots.csv: {
props: {x: ..., "header": ..., ...},
data: "unstructured data (as stored for given extension)",
Expand All @@ -59,30 +59,20 @@ def collect(
from dvc.utils.collections import ensure_list

targets = ensure_list(targets)
data: Dict[str, Dict] = {}
for rev in self.repo.brancher(revs=revs):
# .brancher() adds unwanted workspace
if revs is not None and rev not in revs:
continue
rev = rev or "workspace"
data[rev] = self._collect_from_revision(
revision=rev,
targets=targets,
recursive=recursive,
onerror=onerror,
props=props,
)

errored = errored_revisions(data)
if errored:
from dvc.ui import ui

ui.error_write(
"DVC failed to load some plots for following revisions: "
f"'{', '.join(errored)}'."
)

return data
yield {
rev: self._collect_from_revision(
revision=rev,
targets=targets,
recursive=recursive,
onerror=onerror,
props=props,
)
}

@error_handler
def _collect_from_revision(
Expand Down Expand Up @@ -135,9 +125,22 @@ def show(
if onerror is None:
onerror = onerror_collect

return self.collect(
data: Dict[str, Dict] = {}
for rev_data in self.collect(
targets, revs, recursive, onerror=onerror, props=props
)
):
data.update(rev_data)

errored = errored_revisions(data)
if errored:
from dvc.ui import ui

ui.error_write(
"DVC failed to load some plots for following revisions: "
f"'{', '.join(errored)}'."
)

return data

def diff(self, *args, **kwargs):
from .diff import diff
Expand Down
10 changes: 5 additions & 5 deletions tests/func/plots/test_show.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ def test_plot_wrong_metric_type(tmp_dir, scm, dvc, run_copy_metrics):
)

assert isinstance(
dvc.plots.collect(targets=["metric.txt"], onerror=onerror_collect)[
dvc.plots.show(targets=["metric.txt"], onerror=onerror_collect)[
"workspace"
]["data"]["metric.txt"]["error"],
PlotMetricTypeError,
Expand Down Expand Up @@ -148,7 +148,7 @@ def test_plots_show_overlap(tmp_dir, dvc, run_copy_metrics, clear_before_run):
dvc._reset()

assert isinstance(
dvc.plots.collect(onerror=onerror_collect)["workspace"]["error"],
dvc.plots.show(onerror=onerror_collect)["workspace"]["error"],
OverlappingOutputPathsError,
)

Expand Down Expand Up @@ -183,12 +183,12 @@ def test_dir_plots(tmp_dir, dvc, run_copy_metrics):
assert result["workspace"]["data"][p2]["props"] == props


def test_ignore_binary_file(tmp_dir, dvc, run_copy_metrics):
def test_ignore_parsing_error(tmp_dir, dvc, run_copy_metrics):
with open("file", "wb") as fobj:
fobj.write(b"\xc1")

run_copy_metrics("file", "plot_file.json", plots=["plot_file.json"])
result = dvc.plots.collect(onerror=onerror_collect)
result = dvc.plots.show(onerror=onerror_collect)

assert isinstance(
result["workspace"]["data"]["plot_file.json"]["error"], EncodingError
Expand Down Expand Up @@ -219,7 +219,7 @@ def test_log_errors(
with open(file, "a") as fd:
fd.write("\nMALFORMED!")

result = dvc.plots.collect(onerror=onerror_collect)
result = dvc.plots.show(onerror=onerror_collect)
_, error = capsys.readouterr()

assert isinstance(get_in(result, error_path), YAMLFileCorruptedError)
Expand Down

0 comments on commit db75878

Please sign in to comment.