diff --git a/dvc/command/tag.py b/dvc/command/tag.py index 789dacabd..5701db2e5 100644 --- a/dvc/command/tag.py +++ b/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 diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 48e1309b8..70e1aaa18 100644 --- a/dvc/exceptions.py +++ b/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, ) diff --git a/dvc/repo/metrics/show.py b/dvc/repo/metrics/show.py index e331f096b..8a894d74c 100644 --- a/dvc/repo/metrics/show.py +++ b/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 diff --git a/dvc/stage.py b/dvc/stage.py index 52769264a..115bbd34b 100644 --- a/dvc/stage.py +++ b/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)) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index 5b5f206d8..6fa901bc0 100644 --- a/dvc/utils/__init__.py +++ b/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( diff --git a/dvc/utils/collections.py b/dvc/utils/collections.py index 46d8c9d44..2d7b259fa 100644 --- a/dvc/utils/collections.py +++ b/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__ + ) + ) diff --git a/dvc/utils/compat.py b/dvc/utils/compat.py index b32b38e96..0e104a0f4 100644 --- a/dvc/utils/compat.py +++ b/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 diff --git a/dvc/utils/stage.py b/dvc/utils/stage.py new file mode 100644 index 000000000..e88f3c2b0 --- /dev/null +++ b/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) diff --git a/requirements.txt b/requirements.txt index 5723655b1..324a006be 100644 --- a/requirements.txt +++ b/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" diff --git a/setup.py b/setup.py index afe1a6260..c12d8803f 100644 --- a/setup.py +++ b/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.*", diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..733a1bb2f --- /dev/null +++ b/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 diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 1649db7ed..dfad7599f 100644 --- a/tests/func/test_add.py +++ b/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 diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py index bde04aa4b..1d690c5d1 100644 --- a/tests/func/test_checkout.py +++ b/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) diff --git a/tests/func/test_commit.py b/tests/func/test_commit.py index 6d5ad110d..cb8f00a9d 100644 --- a/tests/func/test_commit.py +++ b/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) diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index b3f49828e..21cc54854 100644 --- a/tests/func/test_data_cloud.py +++ b/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() diff --git a/tests/func/test_move.py b/tests/func/test_move.py index c05d69f15..0ad451115 100644 --- a/tests/func/test_move.py +++ b/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 diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index b86889b42..3a6b7eed2 100644 --- a/tests/func/test_repro.py +++ b/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", "", text) diff --git a/tests/func/test_run.py b/tests/func/test_run.py index c771fa463..9015e4eab 100644 --- a/tests/func/test_run.py +++ b/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): diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index 528f8eddc..1095eefe2 100644 --- a/tests/func/test_stage.py +++ b/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() diff --git a/tests/func/test_tag.py b/tests/func/test_tag.py index cd3ae6042..8230d82d6 100644 --- a/tests/func/test_tag.py +++ b/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) diff --git a/tests/unit/utils/test_collections.py b/tests/unit/utils/test_collections.py new file mode 100644 index 000000000..4dd65abaf --- /dev/null +++ b/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" diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index bb91d0976..3bbe0c3c3 100644 --- a/tests/utils/__init__.py +++ b/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()