From db75878bef2a6c9045fb635be4f0ca1695b8766b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= <pawelredzynski@gmail.com>
Date: Fri, 15 Oct 2021 16:24:02 +0200
Subject: [PATCH] plots: make generator out of collect (#6793)

---
 dvc/repo/plots/__init__.py    | 51 ++++++++++++++++++-----------------
 tests/func/plots/test_show.py | 10 +++----
 2 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py
index e58fc0f79..d2e14896b 100644
--- a/dvc/repo/plots/__init__.py
+++ b/dvc/repo/plots/__init__.py
@@ -3,7 +3,7 @@ import io
 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
 
@@ -47,10 +47,10 @@ class Plots:
         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)",
@@ -59,30 +59,20 @@ class Plots:
         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(
@@ -135,9 +125,22 @@ class Plots:
         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
diff --git a/tests/func/plots/test_show.py b/tests/func/plots/test_show.py
index e1b1c8a23..ede5ce3cf 100644
--- a/tests/func/plots/test_show.py
+++ b/tests/func/plots/test_show.py
@@ -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,
@@ -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,
     )
 
@@ -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
@@ -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)