dvc: retain formatting and comments in .dvc files (#1885)

* test: remove never used copy of load_stage_file()

* dvc: save comments and custom formatting in stage files

* Add a couple of TODOs

* test: test repro doesn't clear comments and formatting in .dvc

As part of this I started writing pytest-like tests and added commonly
used fixtures.

* test: simplify and prettify dvc formatting retained test

Also test that line comments after md5 fields are retained.

* test: fix tests

* dvc: fix stage and metrics opens not closed

* dvc: always use utf-8 in stage files

It used to encode in system default encoding. Not sure if PyYAML was
handling this silently or we were just lucky all the way.

* test: unit test apply_diff()

* test: remove TODOs and ASK

* dvc: polish dvc-formatting PR

* dvc: move stage load/dump utils to a new stage utils file

* dvc: clean up apply_diff()

* test: test stage md5 ignores comments

* test: remove pytests starting marker

dvc/command/tag.py

@@ -1,6 +1,6 @@
-import yaml
 import logging
 
+from dvc.utils import to_yaml_string
 from dvc.exceptions import DvcException
 from dvc.command.base import CmdBase, fix_subparsers, append_doc_link
 
@@ -50,7 +50,7 @@ class CmdTagList(CmdBase):
                     recursive=self.args.recursive,
                 )
                 if tags:
-                    logger.info(yaml.dump(tags, default_flow_style=False))
+                    logger.info(to_yaml_string(tags))
             except DvcException:
                 logger.exception("failed list tags")
                 return 1

dvc/exceptions.py

@@ -214,11 +214,12 @@ class NoMetricsError(DvcException):
 
 
 class StageFileCorruptedError(DvcException):
-    def __init__(self, path):
+    def __init__(self, path, cause=None):
         path = os.path.relpath(path)
         super(StageFileCorruptedError, self).__init__(
             "unable to read stage file: {} "
-            "YAML file structure is corrupted".format(path)
+            "YAML file structure is corrupted".format(path),
+            cause=cause,
         )
 
 

dvc/repo/metrics/show.py

@@ -234,10 +234,14 @@ def _read_metrics(self, metrics, branch):
                 branch=branch,
             )
         else:
-            fd = self.tree.open(out.path)
-            metric = _read_metric(
-                fd, typ=typ, xpath=xpath, rel_path=out.rel_path, branch=branch
-            )
+            with self.tree.open(out.path) as fd:
+                metric = _read_metric(
+                    fd,
+                    typ=typ,
+                    xpath=xpath,
+                    rel_path=out.rel_path,
+                    branch=branch,
+                )
 
         if not metric:
             continue

dvc/stage.py

@@ -1,10 +1,10 @@
 from __future__ import unicode_literals
 
-from dvc.utils.compat import str, open
+from dvc.utils.compat import str
 
+import copy
 import re
 import os
-import yaml
 import subprocess
 import logging
 
@@ -15,7 +15,9 @@ import dvc.prompt as prompt
 import dvc.dependency as dependency
 import dvc.output as output
 from dvc.exceptions import DvcException
-from dvc.utils import dict_md5, fix_env, load_stage_file_fobj
+from dvc.utils import dict_md5, fix_env
+from dvc.utils.collections import apply_diff
+from dvc.utils.stage import load_stage_fd, dump_stage_file
 
 
 logger = logging.getLogger(__name__)
@@ -146,6 +148,7 @@ class Stage(object):
         md5=None,
         locked=False,
         tag=None,
+        state=None,
     ):
         if deps is None:
             deps = []
@@ -161,6 +164,7 @@ class Stage(object):
         self.md5 = md5
         self.locked = locked
         self.tag = tag
+        self._state = state or {}
 
     def __repr__(self):
         return "Stage: '{path}'".format(
@@ -565,7 +569,11 @@ class Stage(object):
         Stage._check_dvc_filename(fname)
         Stage._check_isfile(repo, fname)
 
-        d = load_stage_file_fobj(repo.tree.open(fname), fname)
+        with repo.tree.open(fname) as fd:
+            d = load_stage_fd(fd, fname)
+        # Making a deepcopy since the original structure
+        # looses keys in deps and outs load
+        state = copy.deepcopy(d)
 
         Stage.validate(d, fname=os.path.relpath(fname))
         path = os.path.abspath(fname)
@@ -582,6 +590,7 @@ class Stage(object):
             md5=d.get(Stage.PARAM_MD5),
             locked=d.get(Stage.PARAM_LOCKED, False),
             tag=tag,
+            state=state,
         )
 
         stage.deps = dependency.loadd_from(stage, d.get(Stage.PARAM_DEPS, []))
@@ -618,9 +627,8 @@ class Stage(object):
             )
         )
         d = self.dumpd()
-
-        with open(fname, "w") as fd:
-            yaml.safe_dump(d, fd, default_flow_style=False)
+        apply_diff(d, self._state)
+        dump_stage_file(fname, self._state)
 
         self.repo.scm.track_file(os.path.relpath(fname))
 

dvc/utils/__init__.py

@@ -2,8 +2,7 @@
 
 from __future__ import unicode_literals
 
-import yaml
-from dvc.utils.compat import str, builtin_str, open, cast_bytes_py2
+from dvc.utils.compat import str, builtin_str, open, cast_bytes_py2, StringIO
 
 import os
 import sys
@@ -19,7 +18,7 @@ import colorama
 import re
 import logging
 
-from yaml.scanner import ScannerError
+from ruamel.yaml import YAML
 
 
 logger = logging.getLogger(__name__)
@@ -243,18 +242,16 @@ def current_timestamp():
     return int(nanotime.timestamp(time.time()))
 
 
-def load_stage_file(path):
-    with open(path, "r") as fobj:
-        return load_stage_file_fobj(fobj, path)
+def from_yaml_string(s):
+    return YAML().load(StringIO(s))
 
 
-def load_stage_file_fobj(fobj, path):
-    from dvc.exceptions import StageFileCorruptedError
-
-    try:
-        return yaml.safe_load(fobj) or {}
-    except ScannerError:
-        raise StageFileCorruptedError(path)
+def to_yaml_string(data):
+    stream = StringIO()
+    yaml = YAML()
+    yaml.default_flow_style = False
+    yaml.dump(data, stream)
+    return stream.getvalue()
 
 
 def dvc_walk(

dvc/utils/collections.py

@@ -1,6 +1,47 @@
-from __future__ import unicode_literals
+from __future__ import absolute_import, unicode_literals
+from dvc.utils.compat import Mapping
 
 
 # just simple check for Nones and emtpy strings
 def compact(args):
     return list(filter(bool, args))
+
+
+def apply_diff(src, dest):
+    """Recursively apply changes from stc to dest"""
+    Seq = (list, tuple)
+    Container = (Mapping, list, tuple)
+
+    def is_same_type(a, b):
+        return any(
+            isinstance(a, t) and isinstance(b, t)
+            for t in [str, Mapping, Seq, bool]
+        )
+
+    if isinstance(src, Mapping) and isinstance(dest, Mapping):
+        for key, value in src.items():
+            if isinstance(value, Container) and is_same_type(
+                value, dest.get(key)
+            ):
+                apply_diff(value, dest[key])
+            elif key not in dest or value != dest[key]:
+                dest[key] = value
+        for key in set(dest) - set(src):
+            del dest[key]
+    elif isinstance(src, Seq) and isinstance(dest, Seq):
+        if len(src) != len(dest):
+            dest[:] = src
+        else:
+            for i, value in enumerate(src):
+                if isinstance(value, Container) and is_same_type(
+                    value, dest[i]
+                ):
+                    apply_diff(value, dest[i])
+                elif value != dest[i]:
+                    dest[i] = value
+    else:
+        raise AssertionError(
+            "Can't apply diff from {} to {}".format(
+                src.__class__.__name__, dest.__class__.__name__
+            )
+        )

dvc/utils/compat.py

@@ -1,4 +1,5 @@
 """Handle import compatibility between Python 2 and Python 3"""
+from __future__ import absolute_import
 
 import sys
 import os
@@ -92,12 +93,12 @@ def _makedirs(name, mode=0o777, exist_ok=False):
 
 if is_py2:
     from urlparse import urlparse, urljoin  # noqa: F401
-    from StringIO import StringIO  # noqa: F401
-    from io import BytesIO  # noqa: F401
     from BaseHTTPServer import HTTPServer  # noqa: F401
     from SimpleHTTPServer import SimpleHTTPRequestHandler  # noqa: F401
     import ConfigParser  # noqa: F401
     from io import open  # noqa: F401
+    from pathlib2 import Path  # noqa: F401
+    from collections import Mapping  # noqa: F401
 
     builtin_str = str  # noqa: F821
     bytes = str  # noqa: F821
@@ -109,7 +110,26 @@ if is_py2:
     cast_bytes_py2 = cast_bytes
     makedirs = _makedirs
 
+    import StringIO
+    import io
+
+    class StringIO(StringIO.StringIO):
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            self.close()
+
+    class BytesIO(io.BytesIO):
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            self.close()
+
+
 elif is_py3:
+    from pathlib import Path  # noqa: F401
     from os import makedirs  # noqa: F401
     from urllib.parse import urlparse, urljoin  # noqa: F401
     from io import StringIO, BytesIO  # noqa: F401
@@ -118,6 +138,7 @@ elif is_py3:
         SimpleHTTPRequestHandler,  # noqa: F401
     )  # noqa: F401
     import configparser as ConfigParser  # noqa: F401
+    from collections.abc import Mapping  # noqa: F401
 
     builtin_str = str  # noqa: F821
     str = str  # noqa: F821

dvc/utils/stage.py

@@ -0,0 +1,26 @@
+from dvc.utils.compat import open
+
+from ruamel.yaml.error import YAMLError
+from ruamel.yaml import YAML
+
+from dvc.exceptions import StageFileCorruptedError
+
+
+def load_stage_file(path):
+    with open(path, "r", encoding="utf-8") as fd:
+        return load_stage_fd(fd, path)
+
+
+def load_stage_fd(fd, path):
+    try:
+        yaml = YAML()
+        return yaml.load(fd) or {}
+    except YAMLError as exc:
+        raise StageFileCorruptedError(path, cause=exc)
+
+
+def dump_stage_file(path, data):
+    with open(path, "w", encoding="utf-8") as fd:
+        yaml = YAML()
+        yaml.default_flow_style = False
+        yaml.dump(data, fd)

requirements.txt

@@ -9,7 +9,6 @@ PyInstaller==3.3.1
 colorama>=0.3.9
 configobj>=5.0.6
 networkx>=2.1
-pyyaml>=3.12
 gitpython>=2.1.8
 setuptools>=34.0.0
 nanotime>=0.5.2
@@ -29,3 +28,5 @@ treelib>=1.5.5
 inflect>=2.1.0
 humanize>=0.5.1
 dulwich>=0.19.11
+ruamel.yaml>=0.15.91
+pathlib2==2.3.3; python_version == "2.7"

setup.py

@@ -43,7 +43,6 @@ install_requires = [
     "colorama>=0.3.9",
     "configobj>=5.0.6",
     "networkx>=2.1",
-    "pyyaml>=3.12",
     "gitpython>=2.1.8",
     "setuptools>=34.0.0",
     "nanotime>=0.5.2",
@@ -59,6 +58,7 @@ install_requires = [
     "inflect>=2.1.0",
     "humanize>=0.5.1",
     "dulwich>=0.19.11",
+    "ruamel.yaml==0.15.91",
 ]
 
 # Extra dependencies for remote integrations
@@ -85,7 +85,7 @@ setup(
         "azure": azure,
         "ssh": ssh,
         # NOTE: https://github.com/inveniosoftware/troubleshooting/issues/1
-        ':python_version=="2.7"': ["futures"],
+        ':python_version=="2.7"': ["futures", "pathlib2"],
     },
     keywords="data science, data version control, machine learning",
     python_requires=">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*",

tests/conftest.py

@@ -0,0 +1,52 @@
+import pytest
+from git import Repo
+from git.exc import GitCommandNotFound
+
+from dvc.repo import Repo as DvcRepo
+from .basic_env import TestDirFixture, logger
+
+
+@pytest.fixture(autouse=True)
+def debug():
+    logger.setLevel("DEBUG")
+
+
+# Wrap class like fixture as pytest-like one to avoid code duplication
+@pytest.fixture
+def repo_dir():
+    old_fixture = TestDirFixture()
+    old_fixture.setUp()
+    try:
+        yield old_fixture
+    finally:
+        old_fixture.tearDown()
+
+
+# NOTE: this duplicates code from GitFixture,
+# would fix itself once class-based fixtures are removed
+@pytest.fixture
+def git(repo_dir):
+    # NOTE: handles EAGAIN error on BSD systems (osx in our case).
+    # Otherwise when running tests you might get this exception:
+    #
+    #    GitCommandNotFound: Cmd('git') not found due to:
+    #        OSError('[Errno 35] Resource temporarily unavailable')
+    retries = 5
+    while retries:
+        try:
+            git = Repo.init()
+        except GitCommandNotFound:
+            retries -= 1
+            continue
+        break
+
+    git.index.add([repo_dir.CODE])
+    git.index.commit("add code")
+    return git
+
+
+@pytest.fixture
+def dvc(repo_dir, git):
+    dvc = DvcRepo.init(repo_dir._root_dir)
+    dvc.scm.commit("init dvc")
+    return dvc

tests/func/test_add.py

@@ -12,7 +12,8 @@ from dvc.system import System
 from mock import patch
 
 from dvc.main import main
-from dvc.utils import file_md5, load_stage_file
+from dvc.utils import file_md5
+from dvc.utils.stage import load_stage_file
 from dvc.stage import Stage
 from dvc.exceptions import DvcException, RecursiveAddingWhileUsingFilename
 from dvc.output.base import OutputAlreadyTrackedError

tests/func/test_checkout.py

@@ -2,7 +2,6 @@ import os
 import sys
 import re
 
-import yaml
 import shutil
 import filecmp
 import collections
@@ -12,7 +11,8 @@ from dvc.main import main
 from dvc import progress
 from dvc.repo import Repo as DvcRepo
 from dvc.system import System
-from dvc.utils import walk_files, load_stage_file
+from dvc.utils import walk_files
+from dvc.utils.stage import load_stage_file, dump_stage_file
 from tests.basic_env import TestDvc
 from tests.func.test_repro import TestRepro
 from dvc.stage import Stage, StageFileBadNameError, StageFileDoesNotExistError
@@ -298,14 +298,10 @@ class TestGitIgnoreWhenCheckout(CheckoutBase):
 
 class TestCheckoutMissingMd5InStageFile(TestRepro):
     def test(self):
-        with open(self.file1_stage, "r") as fd:
-            d = yaml.load(fd)
-
+        d = load_stage_file(self.file1_stage)
         del d[Stage.PARAM_OUTS][0][RemoteLOCAL.PARAM_CHECKSUM]
         del d[Stage.PARAM_DEPS][0][RemoteLOCAL.PARAM_CHECKSUM]
-
-        with open(self.file1_stage, "w") as fd:
-            yaml.dump(d, fd)
+        dump_stage_file(self.file1_stage, d)
 
         self.dvc.checkout(force=True)
 

tests/func/test_commit.py

@@ -1,5 +1,4 @@
-import yaml
-from dvc.utils import load_stage_file
+from dvc.utils.stage import load_stage_file, dump_stage_file
 
 from tests.basic_env import TestDvc
 from dvc.stage import StageCommitError
@@ -78,10 +77,8 @@ class TestCommitChangedMd5(TestDvc):
         stage = stages[0]
 
         st = load_stage_file(stage.path)
-
         st["md5"] = "1111111111"
-        with open(stage.path, "w") as fobj:
-            yaml.safe_dump(st, fobj, default_flow_style=False)
+        dump_stage_file(stage.path, st)
 
         with self.assertRaises(StageCommitError):
             self.dvc.commit(stage.path)

tests/func/test_data_cloud.py

@@ -6,7 +6,6 @@ import uuid
 import shutil
 import getpass
 import platform
-import yaml
 import copy
 import logging
 
@@ -27,7 +26,8 @@ from dvc.data_cloud import (
     RemoteHTTP,
 )
 from dvc.remote.base import STATUS_OK, STATUS_NEW, STATUS_DELETED
-from dvc.utils import file_md5, load_stage_file
+from dvc.utils import file_md5
+from dvc.utils.stage import load_stage_file, dump_stage_file
 
 from tests.basic_env import TestDvc
 from tests.utils import spy
@@ -629,8 +629,7 @@ class TestWarnOnOutdatedStage(TestDvc):
         stage_file_path = stage.relpath
         content = load_stage_file(stage_file_path)
         del content["outs"][0]["md5"]
-        with open(stage_file_path, "w") as stage_file:
-            yaml.dump(content, stage_file)
+        dump_stage_file(stage_file_path, content)
 
         with self._caplog.at_level(logging.WARNING, logger="dvc"):
             self._caplog.clear()

tests/func/test_move.py

@@ -3,7 +3,7 @@ import os
 from dvc.main import main
 from dvc.exceptions import DvcException, MoveNotDataSourceError
 from dvc.stage import Stage
-from dvc.utils import load_stage_file
+from dvc.utils.stage import load_stage_file
 
 from tests.basic_env import TestDvc
 from tests.func.test_repro import TestRepro

tests/func/test_repro.py

@@ -1,7 +1,8 @@
-from dvc.utils.compat import str
+from __future__ import unicode_literals
+from dvc.utils.compat import str, urljoin, Path
 
 import os
-import yaml
+import re
 import shutil
 import filecmp
 import getpass
@@ -13,11 +14,12 @@ import uuid
 import paramiko
 from google.cloud import storage as gc
 from flaky.flaky_decorator import flaky
+import pytest
 
 from dvc.main import main
 from dvc.repo import Repo as DvcRepo
 from dvc.utils import file_md5
-from dvc.utils.compat import urljoin
+from dvc.utils.stage import load_stage_file, dump_stage_file
 from dvc.remote.local import RemoteLOCAL
 from dvc.stage import Stage, StageFileDoesNotExistError
 from dvc.system import System
@@ -73,14 +75,12 @@ class TestReproCyclicGraph(TestDvc):
             deps=["bar.txt"], outs=["baz.txt"], cmd="echo baz > baz.txt"
         )
 
-        with open("cycle.dvc", "w") as fd:
-            stage_dump = {
-                "cmd": "echo baz > foo",
-                "deps": [{"path": "baz.txt"}],
-                "outs": [{"path": self.FOO}],
-            }
-
-            yaml.safe_dump(stage_dump, fd, default_flow_style=False)
+        stage_dump = {
+            "cmd": "echo baz > foo",
+            "deps": [{"path": "baz.txt"}],
+            "outs": [{"path": self.FOO}],
+        }
+        dump_stage_file("cycle.dvc", stage_dump)
 
         with self.assertRaises(CyclicGraphError):
             self.dvc.reproduce("cycle.dvc")
@@ -114,14 +114,12 @@ class TestReproWorkingDirectoryAsOutput(TestDvc):
 
         faulty_stage_path = os.path.join("dir2", "something.dvc")
 
-        with open(faulty_stage_path, "w") as fd:
-            output = os.path.join("..", "something")
-            stage_dump = {
-                "cmd": "echo something > {}".format(output),
-                "outs": [{"path": output}],
-            }
-
-            yaml.safe_dump(stage_dump, fd, default_flow_style=False)
+        output = os.path.join("..", "something")
+        stage_dump = {
+            "cmd": "echo something > {}".format(output),
+            "outs": [{"path": output}],
+        }
+        dump_stage_file(faulty_stage_path, stage_dump)
 
         with self.assertRaises(StagePathAsOutputError):
             self.dvc.reproduce(faulty_stage_path)
@@ -156,14 +154,12 @@ class TestReproWorkingDirectoryAsOutput(TestDvc):
 
         error_stage_path = os.path.join(nested_dir, "dir", "error.dvc")
 
-        with open(error_stage_path, "w") as fd:
-            output = os.path.join("..", "..", "something")
-            stage_dump = {
-                "cmd": "echo something > {}".format(output),
-                "outs": [{"path": output}],
-            }
-
-            yaml.safe_dump(stage_dump, fd, default_flow_style=False)
+        output = os.path.join("..", "..", "something")
+        stage_dump = {
+            "cmd": "echo something > {}".format(output),
+            "outs": [{"path": output}],
+        }
+        dump_stage_file(error_stage_path, stage_dump)
 
         # NOTE: os.walk() walks in a sorted order and we need dir2 subdirs to
         # be processed before dir1 to load error.dvc first.
@@ -192,10 +188,8 @@ class TestReproWorkingDirectoryAsOutput(TestDvc):
 
         stage = os.path.join("something-1", "a.dvc")
 
-        with open(stage, "w") as fd:
-            stage_dump = {"cmd": "echo a > a", "outs": [{"path": "a"}]}
-
-            yaml.safe_dump(stage_dump, fd, default_flow_style=False)
+        stage_dump = {"cmd": "echo a > a", "outs": [{"path": "a"}]}
+        dump_stage_file(stage, stage_dump)
 
         try:
             self.dvc.reproduce(stage)
@@ -708,14 +702,10 @@ class TestReproChangedDirData(TestDvc):
 
 class TestReproMissingMd5InStageFile(TestRepro):
     def test(self):
-        with open(self.file1_stage, "r") as fd:
-            d = yaml.load(fd)
-
+        d = load_stage_file(self.file1_stage)
         del d[Stage.PARAM_OUTS][0][RemoteLOCAL.PARAM_CHECKSUM]
         del d[Stage.PARAM_DEPS][0][RemoteLOCAL.PARAM_CHECKSUM]
-
-        with open(self.file1_stage, "w") as fd:
-            yaml.dump(d, fd)
+        dump_stage_file(self.file1_stage, d)
 
         stages = self.dvc.reproduce(self.file1_stage)
         self.assertEqual(len(stages), 1)
@@ -1378,3 +1368,55 @@ class TestShouldDisplayMetricsOnReproWithMetricsOption(TestDvc):
 
         expected_metrics_display = "{}: {}".format(metrics_file, metrics_value)
         self.assertIn(expected_metrics_display, self._caplog.text)
+
+
+@pytest.fixture
+def foo_copy(repo_dir, dvc):
+    stages = dvc.add(repo_dir.FOO)
+    assert len(stages) == 1
+    foo_stage = stages[0]
+    assert foo_stage is not None
+
+    fname = "foo_copy"
+    stage_fname = fname + ".dvc"
+    dvc.run(
+        fname=stage_fname,
+        outs=[fname],
+        deps=[repo_dir.FOO, repo_dir.CODE],
+        cmd="python {} {} {}".format(repo_dir.CODE, repo_dir.FOO, fname),
+    )
+    return {"fname": fname, "stage_fname": stage_fname}
+
+
+def test_dvc_formatting_retained(dvc, foo_copy):
+    root = Path(dvc.root_dir)
+    stage_file = root / foo_copy["stage_fname"]
+
+    # Add comments and custom formatting to stage file
+    lines = list(map(_format_dvc_line, stage_file.read_text().splitlines()))
+    lines.insert(0, "# Starting comment")
+    stage_text = "".join(l + "\n" for l in lines)
+    stage_file.write_text(stage_text)
+
+    # Rewrite data source and repro
+    (root / "foo").write_text("new_foo")
+    dvc.reproduce(foo_copy["stage_fname"])
+
+    # All differences should be only about md5
+    assert _hide_md5(stage_text) == _hide_md5(stage_file.read_text())
+
+
+def _format_dvc_line(line):
+    # Add line comment for all cache and md5 keys
+    if "cache:" in line or "md5:" in line:
+        return line + " # line comment"
+    # Format command as one word per line
+    elif line.startswith("cmd: "):
+        pre, command = line.split(None, 1)
+        return pre + " >\n" + "\n".join("  " + s for s in command.split())
+    else:
+        return line
+
+
+def _hide_md5(text):
+    return re.sub(r"\b[a-f0-9]{32}\b", "<md5>", text)

tests/func/test_run.py

@@ -5,12 +5,12 @@ import mock
 import shutil
 import filecmp
 import subprocess
-import yaml
 
 from dvc.main import main
 from dvc.output import OutputBase
 from dvc.repo import Repo as DvcRepo
-from dvc.utils import file_md5, load_stage_file
+from dvc.utils import file_md5
+from dvc.utils.stage import load_stage_file
 from dvc.system import System
 from dvc.stage import Stage, StagePathNotFoundError, StagePathNotDirectoryError
 from dvc.stage import StageFileBadNameError, MissingDep
@@ -615,15 +615,13 @@ class TestCmdRunWorkingDirectory(TestDvc):
         stage = self.dvc.run(
             cmd="echo test > {}".format(self.FOO), outs=[self.FOO], wdir="."
         )
-        with open(stage.relpath, "r") as fobj:
-            d = yaml.safe_load(fobj)
+        d = load_stage_file(stage.relpath)
         self.assertEqual(d[Stage.PARAM_WDIR], ".")
 
         stage = self.dvc.run(
             cmd="echo test > {}".format(self.BAR), outs=[self.BAR]
         )
-        with open(stage.relpath, "r") as fobj:
-            d = yaml.safe_load(fobj)
+        d = load_stage_file(stage.relpath)
         self.assertEqual(d[Stage.PARAM_WDIR], ".")
 
     def test_fname_changes_path_and_wdir(self):
@@ -640,8 +638,7 @@ class TestCmdRunWorkingDirectory(TestDvc):
         )
 
         # Check that it is dumped properly (relative to fname)
-        with open(stage.relpath, "r") as fobj:
-            d = yaml.safe_load(fobj)
+        d = load_stage_file(stage.relpath)
         self.assertEqual(d[Stage.PARAM_WDIR], "..")
 
     def test_cwd_is_ignored(self):

tests/func/test_stage.py

@@ -1,4 +1,3 @@
-import yaml
 import tempfile
 import os
 
@@ -6,7 +5,7 @@ from dvc.main import main
 from dvc.output.local import OutputLOCAL
 from dvc.remote.local import RemoteLOCAL
 from dvc.stage import Stage, StageFileFormatError
-from dvc.utils import load_stage_file
+from dvc.utils.stage import load_stage_file, dump_stage_file
 
 from tests.basic_env import TestDvc
 
@@ -74,8 +73,6 @@ class TestSchemaDepsOuts(TestSchema):
 
 class TestReload(TestDvc):
     def test(self):
-        import yaml
-
         stages = self.dvc.add(self.FOO)
         self.assertEqual(len(stages), 1)
         stage = stages[0]
@@ -86,9 +83,7 @@ class TestReload(TestDvc):
         # NOTE: checking that reloaded stage didn't change its checksum
         md5 = "11111111111111111111111111111111"
         d[stage.PARAM_MD5] = md5
-
-        with open(stage.relpath, "w") as fobj:
-            yaml.safe_dump(d, fobj, default_flow_style=False)
+        dump_stage_file(stage.relpath, d)
 
         stage = Stage.load(self.dvc, stage.relpath)
         self.assertTrue(stage is not None)
@@ -109,16 +104,13 @@ class TestDefaultWorkingDirectory(TestDvc):
         d = stage.dumpd()
         self.assertEqual(d[stage.PARAM_WDIR], ".")
 
-        with open(stage.relpath, "r") as fobj:
-            d = yaml.safe_load(fobj)
+        d = load_stage_file(stage.relpath)
         self.assertEqual(d[stage.PARAM_WDIR], ".")
 
         del d[stage.PARAM_WDIR]
-        with open(stage.relpath, "w") as fobj:
-            yaml.safe_dump(d, fobj, default_flow_style=False)
+        dump_stage_file(stage.relpath, d)
 
-        with open(stage.relpath, "r") as fobj:
-            d = yaml.safe_load(fobj)
+        d = load_stage_file(stage.relpath)
         self.assertIsNone(d.get(stage.PARAM_WDIR))
 
         with self.dvc.state:
@@ -165,3 +157,13 @@ class TestExternalRemoteResolution(TestDvc):
         assert main(["import", "remote://storage/file", "movie.txt"]) == 0
 
         assert os.path.exists("movie.txt")
+
+
+def test_md5_ignores_comments(repo_dir, dvc):
+    stage, = dvc.add("foo")
+
+    with open(stage.path, "a") as f:
+        f.write("# End comment\n")
+
+    new_stage = Stage.load(dvc, stage.path)
+    assert not new_stage.changed_md5()

tests/func/test_tag.py

@@ -4,6 +4,7 @@ import filecmp
 import logging
 
 from dvc.main import main
+from dvc.utils import from_yaml_string
 
 from tests.basic_env import TestDvc
 
@@ -64,20 +65,19 @@ class TestTagAll(TestDvc):
         ret = main(["tag", "list"])
         self.assertEqual(ret, 0)
 
-        expected = (
-            "bar.dvc:\n"
-            "  bar:\n"
-            "    v1:\n"
-            "      md5: 8978c98bb5a48c2fb5f2c4c905768afa\n"
-            "foo.dvc:\n"
-            "  foo:\n"
-            "    v1:\n"
-            "      md5: acbd18db4cc2f85cedef654fccc4a4d8\n"
-        )
+        expected = {
+            "foo.dvc": {
+                "foo": {"v1": {"md5": "acbd18db4cc2f85cedef654fccc4a4d8"}}
+            },
+            "bar.dvc": {
+                "bar": {"v1": {"md5": "8978c98bb5a48c2fb5f2c4c905768afa"}}
+            },
+        }
 
         stdout = "\n".join(record.message for record in self._caplog.records)
+        received = from_yaml_string(stdout)
 
-        assert expected == stdout
+        assert expected == received
 
         ret = main(["tag", "remove", "v1"])
         self.assertEqual(ret, 0)

tests/unit/utils/test_collections.py

@@ -0,0 +1,49 @@
+from dvc.utils.collections import apply_diff
+
+
+class MyDict(dict):
+    pass
+
+
+class MyInt(int):
+    pass
+
+
+def test_apply_diff_is_inplace():
+    dest = MyDict()
+    dest.attr = 42
+    apply_diff({}, dest)
+
+    assert type(dest) is MyDict, "Preserves class"
+    assert dest.attr == 42, "Preserves custom attrs"
+
+
+def test_apply_diff_mapping():
+    src = {"a": 1}
+    dest = {"b": 2}
+    apply_diff(src, dest)
+    assert dest == src, "Adds and removes keys"
+
+    src = {"a": 1}
+    dest = {"a": MyInt(1)}
+    apply_diff(src, dest)
+    assert type(dest["a"]) is MyInt, "Does not replace equals"
+
+    src = {"d": {"a": 1}}
+    inner = {}
+    dest = {"d": inner}
+    apply_diff(src, dest)
+    assert dest["d"] is inner, "Updates inner dicts"
+
+
+def test_apply_diff_seq():
+    src = [1]
+    dest = [MyInt(1)]
+    apply_diff(src, dest)
+    assert type(dest[0]) is MyInt, "Does not replace equals"
+
+    src = {"l": [1]}
+    inner = []
+    dest = {"l": inner}
+    apply_diff(src, dest)
+    assert dest["l"] is inner, "Updates inner lists"

tests/utils/__init__.py

@@ -1,4 +1,3 @@
-import yaml
 import os
 
 from dvc.scm import Git
@@ -22,11 +21,6 @@ def get_gitignore_content():
         return gitignore.read().splitlines()
 
 
-def load_stage_file(path):
-    with open(path, "r") as fobj:
-        return yaml.safe_load(fobj)
-
-
 @contextmanager
 def cd(newdir):
     prevdir = os.getcwd()