Merge pull request #1669 from pared/rfc_unify_stage_file_loading

stage: extract yaml loading to util method

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)

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 {}

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
         )

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:

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)

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):

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)

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()