Add plots modify (#581)

* add plots modify

* plots modify

* update plots mmodify

* add plots options to zn.plots

* add global option

* extend tests

* extend tests

* test fails

* code for change to outs

* poetry update

* lint

* copy template it not in the cwd

* use outs instead of plots

* fix save

Author: PythonFZ

poetry.lock

@@ -10,7 +10,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.8,<4.0.0"
-dvc = "^2.53.0"
+dvc = "^2.52.0, !=2.53.0"
 pyyaml = "^6.0"
 tqdm = "^4.64.0"
 pandas = "^1.4.3"

pyproject.toml

@@ -0,0 +1,188 @@
+import zntrack
+import pandas as pd
+import pytest
+import pathlib
+import yaml
+from zntrack.utils import run_dvc_cmd
+
+
+class NodeWithPlotsDVC(zntrack.Node):
+    data = zntrack.dvc.plots(
+        zntrack.nwd / "data.csv",
+        x="x",
+        y="y",
+        x_label="x",
+        y_label="y",
+        title="title",
+        template="linear",
+        use_global_plots=False,
+    )
+
+    def run(self) -> None:
+        df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+        df.to_csv(self.data)
+
+    def get_data(self):
+        return pd.read_csv(self.data)
+
+
+class NodeWithPlotsZn(zntrack.Node):
+    data = zntrack.zn.plots(
+        x="x",
+        y="y",
+        x_label="x",
+        y_label="y",
+        title="title",
+        template="linear",
+        use_global_plots=False,
+    )
+
+    def run(self) -> None:
+        self.data = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+
+    def get_data(self):
+        return self.data
+
+
+class NodeWithPlotsZnGlobal(zntrack.Node):
+    data = zntrack.zn.plots(
+        x="x",
+        y=["y", "z"],
+        x_label="x",
+        y_label="y",
+        title="title",
+        template="linear",
+    )
+    metrics = zntrack.zn.metrics()
+
+    def run(self) -> None:
+        self.data = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
+        self.metrics = {"a": 1, "b": 2}
+
+    def get_data(self):
+        return self.data
+
+
+class NodeWithPlotsDVCGlobal(zntrack.Node):
+    data = zntrack.dvc.plots(
+        zntrack.nwd / "data.csv",
+        x="x",
+        y=["y", "z"],
+        x_label="x",
+        y_label="y",
+        title="title",
+        template="linear",
+    )
+    metrics = zntrack.dvc.metrics(zntrack.nwd / "metrics.json")
+
+    def run(self) -> None:
+        df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
+        df.to_csv(self.data)
+        self.metrics.write_text("{'a': 1, 'b': 2}")
+
+    def get_data(self):
+        return pd.read_csv(self.data)
+
+
+class NodeRemoteTemplate(zntrack.Node):
+    data = zntrack.dvc.plots(
+        zntrack.nwd / "data.csv",
+        template=zntrack.__file__,
+    )
+
+    def run(self) -> None:
+        df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
+        df.to_csv(self.data)
+
+
+@pytest.mark.parametrize("cls", [NodeWithPlotsDVC, NodeWithPlotsZn])
+@pytest.mark.parametrize("eager", [True, False])
+def test_NodeWithPlots(proj_path, eager, cls):
+    with zntrack.Project() as project:
+        node = cls()
+    project.run(eager=eager)
+
+    if not eager:
+        node.load()
+        dvc_dict = yaml.safe_load((proj_path / "dvc.yaml").read_text())
+
+        plots = {
+            "x": "x",
+            "y": "y",
+            "x_label": "x",
+            "y_label": "y",
+            "title": "title",
+            "template": "linear",
+        }
+
+        assert (
+            dvc_dict["stages"][node.name]["plots"][0][f"nodes/{node.name}/data.csv"]
+            == plots
+        )
+        run_dvc_cmd(["plots", "show"])
+
+    if isinstance(node, NodeWithPlotsDVC):
+        assert NodeWithPlotsDVC.from_rev().data == pathlib.Path(
+            "nodes", node.name, "data.csv"
+        )
+
+    assert node.get_data()["x"].tolist() == [1, 2, 3]
+
+
+@pytest.mark.parametrize("cls", [NodeWithPlotsZnGlobal, NodeWithPlotsDVCGlobal])
+@pytest.mark.parametrize("eager", [True, False])
+def test_NodeWithPlotsGlobal(proj_path, eager, cls):
+    with zntrack.Project() as project:
+        node = cls()
+    project.run(eager=eager)
+
+    if not eager:
+        node.load()
+        dvc_dict = yaml.safe_load((proj_path / "dvc.yaml").read_text())
+
+        plots = {
+            "x": "x",
+            "y": ["y", "z"],
+            "x_label": "x",
+            "y_label": "y",
+            "title": "title",
+            "template": "linear",
+        }
+
+        assert dvc_dict["plots"][0][f"nodes/{node.name}/data.csv"] == plots
+
+        run_dvc_cmd(["plots", "show"])
+
+    assert node.get_data()["x"].tolist() == [1, 2, 3]
+
+
+def test_multiple_plots_nodes(proj_path):
+    with zntrack.Project(automatic_node_names=True) as project:
+        NodeWithPlotsZnGlobal()
+        NodeWithPlotsZnGlobal()
+        NodeWithPlotsDVCGlobal()
+        a = NodeWithPlotsDVCGlobal()
+    project.run()
+    run_dvc_cmd(["plots", "show"])
+
+    with project:
+        NodeWithPlotsZnGlobal()
+        b = NodeWithPlotsZnGlobal()
+        NodeWithPlotsDVCGlobal()
+        NodeWithPlotsDVCGlobal()
+    project.run()
+    run_dvc_cmd(["plots", "show"])
+    # run_dvc_cmd(["repro", "-f"])
+
+    # check loading
+    NodeWithPlotsDVCGlobal.from_rev(name=a.name)
+    NodeWithPlotsZnGlobal.from_rev(name=b.name)
+
+
+def test_NodeRemoteTemplate(proj_path):
+    with zntrack.Project(automatic_node_names=True) as project:
+        NodeRemoteTemplate()
+        NodeRemoteTemplate()
+    project.run(repro=False)
+
+    assert pathlib.Path("dvc_plots", "templates", "__init__.py").exists()

tests/integration/test_dvc_plots.py

@@ -5,7 +5,7 @@
 
 import znjson
 
-from zntrack.fields.field import Field, FieldGroup
+from zntrack.fields.field import Field, FieldGroup, PlotsMixin
 from zntrack.utils import node_wd
 
 if typing.TYPE_CHECKING:
@@ -130,6 +130,10 @@ def __get__(self, instance: "Node", owner=None):
         return node_wd.ReplaceNWD()(value, nwd=instance.nwd)
 
 
+class PlotsOption(PlotsMixin, DVCOption):
+    """Field with DVC plots kwargs."""
+
+
 def outs(*args, **kwargs) -> DVCOption:
     """Create a outs field."""
     return DVCOption(*args, dvc_option="outs", **kwargs)
@@ -167,9 +171,9 @@ def metrics_no_cache(*args, **kwargs) -> DVCOption:
 
 def plots(*args, **kwargs) -> DVCOption:
     """Create a plots field."""
-    return DVCOption(*args, dvc_option="plots", **kwargs)
+    return PlotsOption(*args, dvc_option="plots", **kwargs)
 
 
 def plots_no_cache(*args, **kwargs) -> DVCOption:
     """Create a plots_no_cache field."""
-    return DVCOption(*args, dvc_option="plots-no-cache", **kwargs)
+    return PlotsOption(*args, dvc_option="plots-no-cache", **kwargs)

zntrack/fields/dvc/__init__.py

@@ -5,8 +5,10 @@
 import json
 import logging
 import pathlib
+import shutil
 import typing
 
+import yaml
 import zninit
 
 from zntrack.utils import LazyOption, config
@@ -196,3 +198,113 @@ def load(self, instance: "Node", lazy: bool = None):
             instance.__dict__[self.name] = LazyOption
         else:
             super().load(instance)
+
+
+class PlotsMixin(Field):
+    """DVC Plots Option including 'dvc plots modify' command."""
+
+    def __init__(
+        self,
+        *args,
+        template=None,
+        x=None,
+        y=None,
+        x_label=None,
+        y_label=None,
+        title=None,
+        use_global_plots: bool = True,
+        **kwargs,
+    ):
+        """Create a DVCOption field.
+
+        Attributes
+        ----------
+        use_global_plots : bool
+            Save the plots config not in 'stages' but in 'plots' in the dvc.yaml file.
+        """
+        super().__init__(*args, **kwargs)
+        self.plots_options = {}
+        self.use_global_plots = use_global_plots
+        if self.use_global_plots:
+            if self.dvc_option == "plots":
+                self.dvc_option = "outs"
+            elif self.dvc_option == "plots-no-cache":
+                self.dvc_option = "outs-no-cache"
+        if template is not None:
+            self.plots_options["--template"] = pathlib.Path(template).as_posix()
+        if x is not None:
+            self.plots_options["-x"] = x
+        if y is not None:
+            self.plots_options["-y"] = y
+        if x_label is not None:
+            self.plots_options["--x-label"] = x_label
+        if y_label is not None:
+            self.plots_options["--y-label"] = y_label
+        if title is not None:
+            self.plots_options["--title"] = title
+
+    def save(self, instance: "Node"):
+        """Save plots options to dvc.yaml, if use_global_plots is True."""
+        if self.plots_options.get("--template") is not None:
+            template = pathlib.Path(self.plots_options["--template"]).resolve()
+            if pathlib.Path.cwd() not in template.parents:
+                # copy template to dvc_plots/templates if it is not in the cwd
+                template_dir = pathlib.Path.cwd() / "dvc_plots" / "templates"
+                template_dir.mkdir(parents=True, exist_ok=True)
+                shutil.copy(template, template_dir)
+                self.plots_options["--template"] = (
+                    (template_dir / template.name)
+                    .relative_to(pathlib.Path.cwd())
+                    .as_posix()
+                )
+
+        with contextlib.suppress(NotImplementedError):
+            super().save(instance=instance)
+        if not self.use_global_plots:
+            return
+
+        dvc_file = pathlib.Path("dvc.yaml")
+        if not dvc_file.exists():
+            dvc_file.write_text(yaml.safe_dump({}))
+        dvc_config = yaml.safe_load(dvc_file.read_text())
+        plots = dvc_config.get("plots", [])
+
+        # remove leading "-/--"
+        for key in list(self.plots_options):
+            if key.startswith("--"):
+                self.plots_options[key[2:]] = self.plots_options[key]
+                del self.plots_options[key]
+            elif key.startswith("-"):
+                self.plots_options[key[1:]] = self.plots_options[key]
+                del self.plots_options[key]
+        # replace "-" with "_"
+        for key in list(self.plots_options):
+            if key.replace("-", "_") != key:
+                self.plots_options[key.replace("-", "_")] = self.plots_options[key]
+                del self.plots_options[key]
+
+        for file in self.get_files(instance):
+            replaced = False
+            for entry in plots:  # entry: dict{filename: {x:, y:, ...}}
+                if pathlib.Path(file) == pathlib.Path(next(iter(entry))):
+                    entry = {pathlib.Path(file).as_posix(): self.plots_options}
+                    replaced = True
+            if not replaced:
+                plots.append({pathlib.Path(file).as_posix(): self.plots_options})
+
+        dvc_config["plots"] = plots
+        dvc_file.write_text(yaml.dump(dvc_config))
+
+    def get_optional_dvc_cmd(self, instance: "Node") -> typing.List[typing.List[str]]:
+        """Add 'dvc plots modify' to this option."""
+        if not self.use_global_plots:
+            cmds = []
+            for file in self.get_files(instance):
+                cmd = ["plots", "modify", pathlib.Path(file).as_posix()]
+                for key, value in self.plots_options.items():
+                    cmd.append(f"{key}")
+                    cmd.append(pathlib.Path(value).as_posix())
+                cmds.append(cmd)
+            return cmds
+        else:
+            return []

zntrack/fields/field.py

@@ -14,7 +14,7 @@
 from znflow import handler
 
 from zntrack import exceptions
-from zntrack.fields.field import DataIsLazyError, Field, FieldGroup, LazyField
+from zntrack.fields.field import DataIsLazyError, Field, FieldGroup, LazyField, PlotsMixin
 from zntrack.utils import module_handler, update_key_val
 
 if typing.TYPE_CHECKING:
@@ -234,7 +234,7 @@ def get_stage_add_argument(self, instance) -> typing.List[tuple]:
         return [(f"--{self.dvc_option}", file.as_posix())]
 
 
-class Plots(LazyField):
+class Plots(PlotsMixin, LazyField):
     """A field that is saved to disk."""
 
     dvc_option: str = "plots"
@@ -246,6 +246,7 @@ def get_files(self, instance) -> list:
 
     def save(self, instance: "Node"):
         """Save the field to disk."""
+        super().save(instance)
         try:
             value = self.get_value_except_lazy(instance)
         except DataIsLazyError: