add util function to group plot definitions

dvc/render/match.py

@@ -8,7 +8,7 @@ import dpath.options
 from funcy import get_in, last
 
 from dvc.repo.plots import _normpath, infer_data_sources
-from dvc.utils.plots import get_plot_id
+from dvc.utils.plots import group_definitions_by_id
 
 from .convert import _get_converter
 
@@ -36,22 +36,9 @@ class PlotsData:
     def group_definitions(self):
         groups = defaultdict(list)
         for rev, rev_content in self.data.items():
-            groups_by_id: Dict = defaultdict(dict)
-            for config_file, config_file_content in (
-                rev_content.get("definitions", {}).get("data", {}).items()
-            ):
-                for plot_id, plot_definition in config_file_content.get(
-                    "data", {}
-                ).items():
-                    groups_by_id[plot_id][config_file] = (rev, plot_id, plot_definition)
-            # only keep config_file if the same plot_id is in multiple config_files
-            for plot_id, configs in groups_by_id.items():
-                if len(configs) == 1:
-                    groups[plot_id].append(next(iter(configs.values())))
-                else:
-                    for config_file, content in configs.items():
-                        full_id = get_plot_id(plot_id, config_file)
-                        groups[full_id].append(content)
+            definitions = rev_content.get("definitions", {}).get("data", {})
+            for plot_id, definition in group_definitions_by_id(definitions).items():
+                groups[plot_id].append((rev, *definition))
         return dict(groups)
 
     def get_definition_data(self, target_files, rev):

dvc/utils/plots.py

@@ -1,4 +1,36 @@
+from collections import defaultdict
+from typing import Dict, Tuple
+
+
 def get_plot_id(config_plot_id: str, config_file_path: str = ""):
     return (
         f"{config_file_path}::{config_plot_id}" if config_file_path else config_plot_id
     )
+
+
+def group_definitions_by_id(
+    definitions: Dict[str, Dict]
+) -> Dict[str, Tuple[str, Dict]]:
+    """
+    Format ID and extracts plot_definition for each plot.
+
+    Arguments:
+        definitions: dict of {config_file: config_file_content}.
+
+    Returns:
+        Dict of {plot_id: (original_plot_id, plot_definition)}.
+    """
+    groups_by_config: Dict = defaultdict(dict)
+    groups_by_id: Dict = {}
+    for config_file, config_file_content in definitions.items():
+        for plot_id, plot_definition in config_file_content.get("data", {}).items():
+            groups_by_config[plot_id][config_file] = (plot_id, plot_definition)
+    # only keep config_file if the same plot_id is in multiple config_files
+    for plot_id, configs in groups_by_config.items():
+        if len(configs) == 1:
+            groups_by_id[plot_id] = next(iter(configs.values()))
+        else:
+            for config_file, content in configs.items():
+                full_id = get_plot_id(plot_id, config_file)
+                groups_by_id[full_id] = content
+    return groups_by_id

tests/unit/utils/test_plots.py

@@ -0,0 +1,18 @@
+from dvc.utils.plots import get_plot_id, group_definitions_by_id
+
+
+def test_get_plot_id():
+    assert get_plot_id("plot_id", "config_path") == "config_path::plot_id"
+    assert get_plot_id("plot_id", "") == "plot_id"
+
+
+def test_group_definitions_by_id():
+    definitions = {
+        "config1": {"data": {"plot1": "definition1", "plot2": "definition2"}},
+        "config2": {"data": {"plot1": "definition1"}},
+    }
+    assert group_definitions_by_id(definitions) == {
+        "config1::plot1": ("plot1", "definition1"),
+        "config2::plot1": ("plot1", "definition1"),
+        "plot2": ("plot2", "definition2"),
+    }