From 20c7b0e3e67a69bfb83b16c5eb784ad35c438e56 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= <pawelredzynski@gmail.com>
Date: Wed, 25 May 2022 14:10:53 +0200
Subject: [PATCH] plots: grouping: stop using dpath.util.search

---
 dvc/render/match.py             | 26 ++++++++++++--------------
 tests/func/test_live.py         |  3 +--
 tests/unit/render/test_match.py | 12 ------------
 tests/unit/test_plots.py        |  6 ++----
 4 files changed, 15 insertions(+), 32 deletions(-)

diff --git a/dvc/render/match.py b/dvc/render/match.py
index 69dff81fc..7f5f623b0 100644
--- a/dvc/render/match.py
+++ b/dvc/render/match.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 import dpath.util
 
@@ -10,20 +10,18 @@ if TYPE_CHECKING:
     from dvc.types import StrPath
 
 
-def get_files(plots_data: Dict) -> List:
-    files = set()
-    for rev in plots_data.keys():
-        for file in plots_data[rev].get("data", {}).keys():
-            files.add(file)
-    sorted_files = sorted(files)
-    return sorted_files
-
-
 def group_by_filename(plots_data: Dict) -> Dict:
-    files = get_files(plots_data)
-    grouped = {
-        file: dpath.util.search(plots_data, ["*", "*", file]) for file in files
-    }
+    grouped: Dict[str, Dict] = {}
+
+    for revision in plots_data.keys():
+        data = plots_data[revision].get("data", {})
+        for file in data.keys():
+            content = data.get(file)
+            if content:
+                dpath.util.new(
+                    grouped, [file, revision, "data", file], content
+                )
+
     return grouped
 
 
diff --git a/tests/func/test_live.py b/tests/func/test_live.py
index 6763431f6..8bc31a9c6 100644
--- a/tests/func/test_live.py
+++ b/tests/func/test_live.py
@@ -6,7 +6,6 @@ import pytest
 from funcy import first
 
 from dvc import stage as stage_module
-from dvc.render.match import get_files
 
 pytest.importorskip("dvclive", reason="no dvclive")
 
@@ -122,7 +121,7 @@ def test_live_provides_metrics(tmp_dir, dvc, live_stage):
 
     assert (tmp_dir / "logs").is_dir()
     plots_data = dvc.plots.show()
-    files = get_files(plots_data)
+    files = list(plots_data["workspace"]["data"])
     assert os.path.join("logs", "scalars", "accuracy.tsv") in files
     assert os.path.join("logs", "scalars", "loss.tsv") in files
     assert os.path.join("logs", "images", "0", "image.jpg") in files
diff --git a/tests/unit/render/test_match.py b/tests/unit/render/test_match.py
index f9b9aad81..53ca64540 100644
--- a/tests/unit/render/test_match.py
+++ b/tests/unit/render/test_match.py
@@ -1,22 +1,10 @@
 from dvc.render.match import (
-    get_files,
     group_by_filename,
     match_renderers,
     squash_plots_properties,
 )
 
 
-def test_get_files_on_error():
-    plots_data = {
-        "v2": {"error": "ERROR"},
-        "v1": {
-            "data": {"file.json": {"data": [{"y": 2}, {"y": 3}], "props": {}}}
-        },
-    }
-    files = get_files(plots_data)
-    assert files == ["file.json"]
-
-
 def test_group_by_filename():
     error = FileNotFoundError()
     data = {
diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py
index 7bbc4a77f..9b9136c72 100644
--- a/tests/unit/test_plots.py
+++ b/tests/unit/test_plots.py
@@ -1,8 +1,6 @@
 import json
 import os
 
-from dvc.render.match import get_files
-
 
 def test_plots_order(tmp_dir, dvc):
     tmp_dir.gen(
@@ -28,9 +26,9 @@ def test_plots_order(tmp_dir, dvc):
             name="stage2",
         )
 
-    assert get_files(dvc.plots.show()) == [
+    assert list(dvc.plots.show()["workspace"]["data"]) == [
         "p.json",
+        os.path.join("sub", "p4.json"),
         "p1.json",
         os.path.join("sub", "p3.json"),
-        os.path.join("sub", "p4.json"),
     ]