dvc: preserve comments/meta on add/imp/run (#5001)

dvc/dvcfile.py

@@ -223,8 +223,6 @@ class PipelineFile(FileMixin):
 
             if existing_entry:
                 orig_stage_data = data["stages"][stage.name]
-                if "meta" in orig_stage_data:
-                    stage_data[stage.name]["meta"] = orig_stage_data["meta"]
                 apply_diff(stage_data[stage.name], orig_stage_data)
             else:
                 data["stages"].update(stage_data)

dvc/repo/add.py

@@ -154,7 +154,7 @@ def _create_stages(
 ):
     from glob import iglob
 
-    from dvc.stage import Stage, create_stage
+    from dvc.stage import Stage, create_stage, restore_meta
 
     if glob:
         expanded_targets = [
@@ -181,6 +181,7 @@ def _create_stages(
             outs=[out],
             external=external,
         )
+        restore_meta(stage)
         if stage.can_be_skipped:
             stage = None
         else:

dvc/repo/imp_url.py

@@ -21,7 +21,7 @@ def imp_url(
     desc=None,
 ):
     from dvc.dvcfile import Dvcfile
-    from dvc.stage import Stage, create_stage
+    from dvc.stage import Stage, create_stage, restore_meta
 
     out = resolve_output(url, out)
     path, wdir, out = resolve_paths(self, out)
@@ -43,7 +43,7 @@ def imp_url(
         outs=[out],
         erepo=erepo,
     )
-
+    restore_meta(stage)
     if stage.can_be_skipped:
         return None
 

dvc/repo/run.py

@@ -78,7 +78,7 @@ def _check_stage_exists(dvcfile, stage):
 @scm_context
 def run(self, fname=None, no_exec=False, single_stage=False, **kwargs):
     from dvc.dvcfile import PIPELINE_FILE, Dvcfile
-    from dvc.stage import Stage, create_stage
+    from dvc.stage import Stage, create_stage, restore_meta
 
     if not kwargs.get("cmd"):
         raise InvalidArgumentError("command is not specified")
@@ -113,6 +113,7 @@ def run(self, fname=None, no_exec=False, single_stage=False, **kwargs):
     stage = create_stage(
         stage_cls, repo=self, path=path, params=params, **kwargs
     )
+    restore_meta(stage)
     if kwargs.get("run_cache", True) and stage.can_be_skipped:
         return None
 

dvc/stage/__init__.py

@@ -55,6 +55,7 @@ def loads_from(cls, repo, path, wdir, data):
                 Stage.PARAM_ALWAYS_CHANGED,
                 Stage.PARAM_MD5,
                 Stage.PARAM_DESC,
+                Stage.PARAM_META,
                 "name",
             ],
         ),
@@ -90,6 +91,24 @@ def create_stage(cls, repo, path, external=False, **kwargs):
     return stage
 
 
+def restore_meta(stage):
+    from .exceptions import StageNotFound
+
+    if not stage.dvcfile.exists():
+        return
+
+    try:
+        old = stage.reload()
+    except StageNotFound:
+        return
+
+    # will be used to restore comments later
+    # noqa, pylint: disable=protected-access
+    stage._stage_text = old._stage_text
+
+    stage.meta = old.meta
+
+
 class Stage(params.StageParams):
     # pylint:disable=no-value-for-parameter
     # rwlocked() confuses pylint
@@ -109,6 +128,7 @@ class Stage(params.StageParams):
         stage_text=None,
         dvcfile=None,
         desc=None,
+        meta=None,
     ):
         if deps is None:
             deps = []
@@ -127,6 +147,7 @@ class Stage(params.StageParams):
         self._stage_text = stage_text
         self._dvcfile = dvcfile
         self.desc = desc
+        self.meta = meta
         self.raw_data = RawData()
 
     @property

dvc/stage/loader.py

@@ -76,6 +76,7 @@ class StageLoader(Mapping):
         stage = loads_from(PipelineStage, dvcfile.repo, path, wdir, stage_data)
         stage.name = name
         stage.desc = stage_data.get(Stage.PARAM_DESC)
+        stage.meta = stage_data.get(Stage.PARAM_META)
 
         deps = project(stage_data, [stage.PARAM_DEPS, stage.PARAM_PARAMS])
         fill_stage_dependencies(stage, **deps)

dvc/stage/serialize.py

@@ -134,6 +134,7 @@ def to_pipeline_file(stage: "PipelineStage"):
         (stage.PARAM_PLOTS, plots),
         (stage.PARAM_FROZEN, stage.frozen),
         (stage.PARAM_ALWAYS_CHANGED, stage.always_changed),
+        (stage.PARAM_META, stage.meta),
     ]
     return {
         stage.name: OrderedDict([(key, value) for key, value in res if value])
@@ -185,10 +186,6 @@ def to_single_stage_file(stage: "Stage"):
     text = stage._stage_text  # noqa, pylint: disable=protected-access
     if text is not None:
         saved_state = parse_yaml_for_update(text, stage.path)
-        # Stage doesn't work with meta in any way, so .dumpd() doesn't
-        # have it. We simply copy it over.
-        if "meta" in saved_state:
-            state["meta"] = saved_state["meta"]
         apply_diff(state, saved_state)
         state = saved_state
     return state

dvc/stage/utils.py

@@ -207,6 +207,7 @@ def get_dump(stage):
             stage.PARAM_DEPS: [d.dumpd() for d in stage.deps],
             stage.PARAM_OUTS: [o.dumpd() for o in stage.outs],
             stage.PARAM_ALWAYS_CHANGED: stage.always_changed,
+            stage.PARAM_META: stage.meta,
         }.items()
         if value
     }

tests/func/test_add.py

@@ -2,6 +2,7 @@ import filecmp
 import logging
 import os
 import shutil
+import textwrap
 import time
 
 import colorama
@@ -922,3 +923,27 @@ def test_add_with_cache_link_error(tmp_dir, dvc, mocker, caplog):
     assert (tmp_dir / ".dvc" / "cache").read_text() == {
         "ac": {"bd18db4cc2f85cedef654fccc4a4d8": "foo"}
     }
+
+
+def test_add_preserve_meta(tmp_dir, dvc):
+    text = textwrap.dedent(
+        """\
+        # top comment
+        outs:
+        - path: foo # out comment
+        meta: some metadata
+    """
+    )
+    tmp_dir.gen("foo.dvc", text)
+
+    tmp_dir.dvc_gen("foo", "foo")
+    assert (tmp_dir / "foo.dvc").read_text() == textwrap.dedent(
+        """\
+        # top comment
+        outs:
+        - path: foo # out comment
+          md5: acbd18db4cc2f85cedef654fccc4a4d8
+          size: 3
+        meta: some metadata
+    """
+    )

tests/func/test_dvcfile.py

@@ -315,8 +315,8 @@ def test_dvcfile_dump_preserves_meta(tmp_dir, dvc, run_copy):
 
     data = dvcfile._load()[0]
     metadata = {"name": "copy-file"}
+    stage.meta = metadata
     data["stages"]["run_copy"]["meta"] = metadata
-    dump_yaml(dvcfile.path, data)
 
     dvcfile.dump(stage)
     assert dvcfile._load()[0] == data

tests/func/test_import_url.py

@@ -1,4 +1,5 @@
 import os
+import textwrap
 from uuid import uuid4
 
 import pytest
@@ -207,3 +208,36 @@ def test_import_url_dir(tmp_dir, dvc, workspace, stage_md5, dir_md5):
     )
 
     assert dvc.status() == {}
+
+
+def test_import_url_preserve_meta(tmp_dir, dvc):
+    text = textwrap.dedent(
+        """\
+        # top comment
+        deps:
+        - path: foo # dep comment
+        outs:
+        - path: bar # out comment
+        meta: some metadata
+    """
+    )
+    tmp_dir.gen("bar.dvc", text)
+
+    tmp_dir.gen("foo", "foo")
+    dvc.imp_url("foo", out="bar")
+    assert (tmp_dir / "bar.dvc").read_text() == textwrap.dedent(
+        """\
+        # top comment
+        deps:
+        - path: foo # dep comment
+          md5: acbd18db4cc2f85cedef654fccc4a4d8
+          size: 3
+        outs:
+        - path: bar # out comment
+          md5: acbd18db4cc2f85cedef654fccc4a4d8
+          size: 3
+        meta: some metadata
+        md5: be7ade0aa89cc8d56e320867a9de9740
+        frozen: true
+    """
+    )

tests/func/test_run_single_stage.py

@@ -1005,10 +1005,7 @@ def test_metrics_dir(tmp_dir, dvc, caplog, run_copy_metrics, metrics_type):
     )
 
 
-def test_run_force_doesnot_preserve_comments_and_meta(tmp_dir, dvc, run_copy):
-    """Depends on loading of stage on `run` where we don't check the file
-    for stage already exists, so we don't copy `stage_text` over due to which
-    `meta` and `comments` don't get preserved."""
+def test_run_force_preserves_comments_and_meta(tmp_dir, dvc, run_copy):
     tmp_dir.gen({"foo": "foo", "foo1": "foo1"})
     text = textwrap.dedent(
         """\
@@ -1017,7 +1014,7 @@ def test_run_force_doesnot_preserve_comments_and_meta(tmp_dir, dvc, run_copy):
       - path: copy.py
       - path: foo
       outs:
-      # comment not preserved
+      # comment preserved
       - path: bar
       meta:
         name: copy-foo-bar
@@ -1030,5 +1027,5 @@ def test_run_force_doesnot_preserve_comments_and_meta(tmp_dir, dvc, run_copy):
 
     run_copy("foo1", "bar1", single_stage=True, force=True, fname="bar.dvc")
 
-    assert "comment" not in (tmp_dir / "bar.dvc").read_text()
-    assert "meta" not in (tmp_dir / "bar.dvc").read_text()
+    assert "comment" in (tmp_dir / "bar.dvc").read_text()
+    assert "meta" in (tmp_dir / "bar.dvc").read_text()