support lists for log_plot y val (#837)

Author: Dave Berenbaum

src/dvclive/live.py

@@ -562,7 +562,7 @@ def log_plot(
         name: str,
         datapoints: Union[pd.DataFrame, np.ndarray, List[Dict]],
         x: str,
-        y: str,
+        y: Union[str, list[str]],
         template: Optional[str] = "linear",
         title: Optional[str] = None,
         x_label: Optional[str] = None,
@@ -579,7 +579,8 @@ def log_plot(
             datapoints (pd.DataFrame | np.ndarray | List[Dict]): Pandas DataFrame, Numpy
                 Array or List of dictionaries containing the data for the plot.
             x (str): name of the key (present in the dictionaries) to use as the x axis.
-            y (str): name of the key (present in the dictionaries) to use the y axis.
+            y (str | list[str]): name of the key or keys (present in the
+                dictionaries) to use the y axis.
             template (str): name of the `DVC plots template` to use. Defaults to
                 `"linear"`.
             title (str): title to be displayed. Defaults to

src/dvclive/plots/custom.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
 from dvclive.serialize import dump_json
 
@@ -15,7 +15,7 @@ def __init__(
         name: str,
         output_folder: str,
         x: str,
-        y: str,
+        y: Union[str, list[str]],
         template: Optional[str],
         title: Optional[str] = None,
         x_label: Optional[str] = None,

tests/plots/test_custom.py

@@ -29,3 +29,30 @@ def test_log_custom_plot(tmp_dir):
         "x_label": "x_label",
         "y_label": "y_label",
     }
+
+
+def test_log_custom_plot_multi_y(tmp_dir):
+    live = Live()
+    out = tmp_dir / live.plots_dir / CustomPlot.subfolder
+
+    datapoints = [{"x": 1, "y1": 2, "y2": 3}, {"x": 4, "y1": 5, "y2": 6}]
+    live.log_plot(
+        "custom_linear",
+        datapoints,
+        x="x",
+        y=["y1", "y2"],
+        template="linear",
+        title="custom_title",
+        x_label="x_label",
+        y_label="y_label",
+    )
+
+    assert json.loads((out / "custom_linear.json").read_text()) == datapoints
+    assert live._plots["custom_linear"].plot_config == {
+        "template": "linear",
+        "title": "custom_title",
+        "x": "x",
+        "y": ["y1", "y2"],
+        "x_label": "x_label",
+        "y_label": "y_label",
+    }