plots: accept any number of y data sources

demo125 · Dec 20, 2022 · 6fc0a6f · 6fc0a6f
1 parent b4b5db4
commit 6fc0a6f
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 2 deletions.
diff --git a/dvc/render/converter/vega.py b/dvc/render/converter/vega.py
@@ -134,8 +134,12 @@ def _infer_x_y(self):
             self.inferred_properties["x"] = {}
             # If multiple y files, duplicate x for each file.
             if isinstance(y, dict):
-                for file in y.keys():
-                    self.inferred_properties["x"][file] = x
+                for file, fields in y.items():
+                    # Duplicate x for each y.
+                    if isinstance(fields, list):
+                        self.inferred_properties["x"][file] = [x] * len(fields)
+                    else:
+                        self.inferred_properties["x"][file] = x
             # Otherwise use plot ID as file.
             else:
                 self.inferred_properties["x"][self.plot_id] = x

diff --git a/tests/unit/render/test_vega_converter.py b/tests/unit/render/test_vega_converter.py
@@ -331,6 +331,54 @@ def test_finding_lists(dictionary, expected_result):
             },
             id="y_list",
         ),
+        pytest.param(
+            {
+                "f": {"metric": [{"v": 1, "v2": 0.1, "v3": 0.01}]},
+                "f2": {"metric": [{"v": 1, "v2": 0.1}]},
+            },
+            {"y": {"f": ["v2", "v3"], "f2": ["v2"]}, "x": "v"},
+            [
+                {
+                    "dvc_inferred_y_value": 0.1,
+                    "v": 1,
+                    "v2": 0.1,
+                    "v3": 0.01,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v2",
+                    },
+                },
+                {
+                    "dvc_inferred_y_value": 0.01,
+                    "v": 1,
+                    "v2": 0.1,
+                    "v3": 0.01,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v3",
+                    },
+                },
+                {
+                    "dvc_inferred_y_value": 0.1,
+                    "v": 1,
+                    "v2": 0.1,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f2",
+                        "field": "v2",
+                    },
+                },
+            ],
+            {
+                "x": "v",
+                "y": "dvc_inferred_y_value",
+                "x_label": "v",
+                "y_label": "y",
+            },
+            id="multi_source_y_single_x",
+        ),
     ],
 )
 def test_convert(