From 8ee53f8c81782f546a1736a62af218f943bfcb05 Mon Sep 17 00:00:00 2001 From: pared Date: Thu, 28 Feb 2019 16:27:16 +0100 Subject: [PATCH] stage: extract yaml loading to util method --- dvc/stage.py | 5 ++--- dvc/utils/__init__.py | 7 +++++++ tests/test_add.py | 16 ++++++---------- tests/test_commit.py | 5 ++--- tests/test_data_cloud.py | 5 ++--- tests/test_move.py | 3 ++- tests/test_stage.py | 8 +++----- tests/utils/__init__.py | 6 ------ 8 files changed, 24 insertions(+), 31 deletions(-) diff --git a/dvc/stage.py b/dvc/stage.py index ce3d69414..e8328b2bf 100644 --- a/dvc/stage.py +++ b/dvc/stage.py @@ -14,7 +14,7 @@ import dvc.logger as logger import dvc.dependency as dependency import dvc.output as output from dvc.exceptions import DvcException -from dvc.utils import dict_md5, fix_env +from dvc.utils import dict_md5, fix_env, load_stage_file class StageCmdFailedError(DvcException): @@ -488,8 +488,7 @@ class Stage(object): if not Stage.is_stage_file(fname): raise StageFileIsNotDvcFileError(fname) - with open(fname, "r") as fd: - d = yaml.safe_load(fd) or {} + d = load_stage_file(fname) Stage.validate(d, fname=os.path.relpath(fname)) path = os.path.abspath(fname) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index d6789d8d9..13c501c3b 100644 --- a/dvc/utils/__init__.py +++ b/dvc/utils/__init__.py @@ -1,6 +1,8 @@ """Helpers for other modules.""" from __future__ import unicode_literals + +import yaml from dvc.utils.compat import str, builtin_str, open, cast_bytes_py2 import os @@ -225,3 +227,8 @@ def tmp_fname(fname): def current_timestamp(): return int(nanotime.timestamp(time.time())) + + +def load_stage_file(path): + with open(path, "r") as fobj: + return yaml.safe_load(fobj) or {} diff --git a/tests/test_add.py b/tests/test_add.py index f97ce8387..89fead9d8 100644 --- a/tests/test_add.py +++ b/tests/test_add.py @@ -3,7 +3,6 @@ from __future__ import unicode_literals import os import dvc -import yaml import time import shutil import filecmp @@ -13,7 +12,7 @@ from dvc.system import System from mock import patch from dvc.main import main -from dvc.utils import file_md5 +from dvc.utils import file_md5, load_stage_file from dvc.stage import Stage from dvc.exceptions import DvcException from dvc.output.base import OutputAlreadyTrackedError @@ -185,17 +184,15 @@ class TestAddLocalRemoteFile(TestDvc): ret = main(["add", foo]) self.assertEqual(ret, 0) - with open("foo.dvc", "r") as fobj: - d = yaml.safe_load(fobj) - self.assertEqual(d["outs"][0]["path"], foo) + d = load_stage_file("foo.dvc") + self.assertEqual(d["outs"][0]["path"], foo) bar = os.path.join(cwd, self.BAR) ret = main(["add", bar]) self.assertEqual(ret, 0) - with open("bar.dvc", "r") as fobj: - d = yaml.safe_load(fobj) - self.assertEqual(d["outs"][0]["path"], bar) + d = load_stage_file("bar.dvc") + self.assertEqual(d["outs"][0]["path"], bar) class TestCmdAdd(TestDvc): @@ -359,8 +356,7 @@ class SymlinkAddTestBase(TestDvc): stage_file = self.data_file_name + Stage.STAGE_FILE_SUFFIX self.assertTrue(os.path.exists(stage_file)) - with open(stage_file, "r") as fobj: - d = yaml.safe_load(fobj) + d = load_stage_file(stage_file) relative_data_path = posixpath.join( self.link_name, self.data_file_name ) diff --git a/tests/test_commit.py b/tests/test_commit.py index c71dfd025..6d5ad110d 100644 --- a/tests/test_commit.py +++ b/tests/test_commit.py @@ -1,4 +1,5 @@ import yaml +from dvc.utils import load_stage_file from tests.basic_env import TestDvc from dvc.stage import StageCommitError @@ -76,9 +77,7 @@ class TestCommitChangedMd5(TestDvc): self.assertEqual(len(stages), 1) stage = stages[0] - st = None - with open(stage.path, "r") as fobj: - st = yaml.safe_load(fobj) + st = load_stage_file(stage.path) st["md5"] = "1111111111" with open(stage.path, "w") as fobj: diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index fb6490e7c..6592e833d 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -28,7 +28,7 @@ from dvc.data_cloud import ( RemoteHTTP, ) from dvc.remote.base import STATUS_OK, STATUS_NEW, STATUS_DELETED -from dvc.utils import file_md5 +from dvc.utils import file_md5, load_stage_file from tests.basic_env import TestDvc from tests.utils import spy @@ -632,8 +632,7 @@ class TestWarnOnOutdatedStage(TestDvc): self.main(["push"]) stage_file_path = stage.relpath - with open(stage_file_path, "r") as stage_file: - content = yaml.safe_load(stage_file) + 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) diff --git a/tests/test_move.py b/tests/test_move.py index 297400ae1..c135cf16a 100644 --- a/tests/test_move.py +++ b/tests/test_move.py @@ -3,10 +3,11 @@ 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 tests.basic_env import TestDvc from tests.test_repro import TestRepro -from tests.utils import load_stage_file, cd +from tests.utils import cd class TestMove(TestDvc): diff --git a/tests/test_stage.py b/tests/test_stage.py index e44595c8d..9f79268d1 100644 --- a/tests/test_stage.py +++ b/tests/test_stage.py @@ -1,6 +1,7 @@ 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 tests.basic_env import TestDvc @@ -75,8 +76,7 @@ class TestReload(TestDvc): stage = stages[0] self.assertTrue(stage is not None) - with open(stage.relpath, "r") as fobj: - d = yaml.safe_load(fobj) + d = load_stage_file(stage.relpath) # NOTE: checking that reloaded stage didn't change its checksum md5 = "11111111111111111111111111111111" @@ -89,7 +89,5 @@ class TestReload(TestDvc): self.assertTrue(stage is not None) stage.dump() - with open(stage.relpath, "r") as fobj: - d = yaml.safe_load(fobj) - + d = load_stage_file(stage.relpath) self.assertEqual(d[stage.PARAM_MD5], md5) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index d39ef5d9a..8a9f854a7 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -1,4 +1,3 @@ -import yaml import os from mock import MagicMock from contextlib import contextmanager @@ -15,11 +14,6 @@ def spy(method_to_decorate): return wrapper -def load_stage_file(path): - with open(path, "r") as fobj: - return yaml.safe_load(fobj) - - @contextmanager def cd(newdir): prevdir = os.getcwd()