dvc: bring back dvc lock with dvc unlock

Fixes #523

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/cli.py

@@ -22,6 +22,7 @@ from dvc.command.remote import CmdRemoteAdd, CmdRemoteRemove, CmdRemoteModify, C
 from dvc.command.metrics import CmdMetricsShow, CmdMetricsAdd, CmdMetricsRemove
 from dvc.command.install import CmdInstall
 from dvc.command.root import CmdRoot
+from dvc.command.lock import CmdLock, CmdUnlock
 from dvc.stage import Stage
 from dvc import VERSION
 
@@ -498,6 +499,27 @@ def parse_args(argv=None):
                         help='Relative path to project\'s directory')
     root_parser.set_defaults(func=CmdRoot)
 
+    # Lock
+    lock_parser = subparsers.add_parser(
+                        'lock',
+                        parents=[parent_parser],
+                        help='Lock DVC file')
+    lock_parser.add_argument('targets',
+                        nargs='+',
+                        help='DVC files.')
+    lock_parser.set_defaults(func=CmdLock)
+
+    # Unlock
+    unlock_parser = subparsers.add_parser(
+                        'unlock',
+                        parents=[parent_parser],
+                        help='Unlock DVC file')
+    unlock_parser.add_argument('targets',
+                        nargs='+',
+                        help='DVC files.')
+    unlock_parser.set_defaults(func=CmdUnlock)
+
+
     args = parser.parse_args(argv)
 
     return args

dvc/command/lock.py

@@ -0,0 +1,24 @@
+from dvc.exceptions import DvcException
+from dvc.command.common.base import CmdBase
+
+
+class CmdLockBase(CmdBase):
+    def run(self, unlock=False):
+        for target in self.args.targets:
+            try:
+                self.project.lock_stage(target, unlock=unlock)
+            except DvcException as ex:
+                msg = 'Failed to {}lock \'{}\''.format('un' if unlock else '', target)
+                self.project.logger.error(msg, ex)
+                return 1
+        return 0
+
+
+class CmdLock(CmdLockBase):
+    def run(self):
+        return super(CmdLock, self).run(False)
+
+
+class CmdUnlock(CmdLockBase):
+    def run(self):
+        return super(CmdUnlock, self).run(True)

dvc/project.py

@@ -138,6 +138,16 @@ class Project(object):
 
         return stage
 
+    def lock_stage(self, target, unlock=False):
+        if not Stage.is_stage_file(target):
+            raise NotDvcFileError(target)
+
+        stage = Stage.load(self, target)
+        stage.locked = False if unlock else True
+        stage.dump()
+
+        return stage
+
     def move(self, from_path, to_path):
         from_out = Output.loads_from(Stage(self, cwd=os.curdir), [from_path])[0]
 
@@ -431,6 +441,8 @@ class Project(object):
         for stage in stages:
             node = os.path.relpath(stage.path, self.root_dir)
             G.add_node(node, stage=stage)
+            if stage.locked:
+                continue
             for dep in stage.deps:
                 for out in outs:
                     if out.path != dep.path and not dep.path.startswith(out.path + out.sep):

dvc/stage.py

@@ -42,15 +42,17 @@ class Stage(object):
     PARAM_CMD = 'cmd'
     PARAM_DEPS = 'deps'
     PARAM_OUTS = 'outs'
+    PARAM_LOCKED = 'locked'
 
     SCHEMA = {
         schema.Optional(PARAM_MD5): schema.Or(str, None),
         schema.Optional(PARAM_CMD): schema.Or(str, None),
         schema.Optional(PARAM_DEPS): schema.Or(schema.And(list, schema.Schema([dependency.SCHEMA])), None),
         schema.Optional(PARAM_OUTS): schema.Or(schema.And(list, schema.Schema([output.SCHEMA])), None),
+        schema.Optional(PARAM_LOCKED): bool,
     }
 
-    def __init__(self, project, path=None, cmd=None, cwd=None, deps=[], outs=[], md5=None):
+    def __init__(self, project, path=None, cmd=None, cwd=None, deps=[], outs=[], md5=None, locked=False):
         self.project = project
         self.path = path
         self.cmd = cmd
@@ -58,6 +60,7 @@ class Stage(object):
         self.outs = outs
         self.deps = deps
         self.md5 = md5
+        self.locked = locked
 
     @property
     def relpath(self):
@@ -108,7 +111,12 @@ class Stage(object):
         if self.is_callback:
             ret = True
 
-        for entry in itertools.chain(self.outs, self.deps):
+        if self.locked:
+            entries = self.outs
+        else:
+            entries = itertools.chain(self.outs, self.deps)
+
+        for entry in entries:
             if entry.changed():
                 ret = True
 
@@ -134,7 +142,7 @@ class Stage(object):
         if not self.changed() and not force:
             return None
 
-        if self.cmd or self.is_import:
+        if (self.cmd or self.is_import) and not self.locked:
             # Removing outputs only if we actually have command to reproduce
             self.remove_outs(ignore_remove=False)
 
@@ -162,12 +170,14 @@ class Stage(object):
         cwd = os.path.dirname(path)
         cmd = d.get(Stage.PARAM_CMD, None)
         md5 = d.get(Stage.PARAM_MD5, None)
+        locked = d.get(Stage.PARAM_LOCKED, False)
 
         stage = Stage(project=project,
                       path=path,
                       cmd=cmd,
                       cwd=cwd,
-                      md5=md5)
+                      md5=md5,
+                      locked=locked)
 
         stage.deps = dependency.loadd_from(stage, d.get(Stage.PARAM_DEPS, []))
         stage.outs = output.loadd_from(stage, d.get(Stage.PARAM_OUTS, []))
@@ -182,14 +192,16 @@ class Stage(object):
               outs_no_cache=[],
               metrics_no_cache=[],
               fname=None,
-              cwd=os.curdir):
+              cwd=os.curdir,
+              locked=False):
         cwd = os.path.abspath(cwd)
         path = os.path.join(cwd, fname)
 
         stage = Stage(project=project,
                       path=path,
                       cmd=cmd,
-                      cwd=cwd)
+                      cwd=cwd,
+                      locked=locked)
 
         stage.outs = output.loads_from(stage, outs, use_cache=True)
         stage.outs += output.loads_from(stage, outs_no_cache, use_cache=False)
@@ -217,6 +229,9 @@ class Stage(object):
         if len(outs):
             ret[Stage.PARAM_OUTS] = outs
 
+        if self.locked:
+            ret[Stage.PARAM_LOCKED] = self.locked
+
         ret[Stage.PARAM_MD5] = dict_md5(ret)
 
         return ret
@@ -236,7 +251,10 @@ class Stage(object):
             out.save()
 
     def run(self):
-        if self.is_import:
+        if self.locked:
+            self.project.logger.info(u'Verifying outputs in locked stage \'{}\''.format(self.relpath))
+            self.check_missing_outputs()
+        elif self.is_import:
             msg = u'Importing \'{}\' -> \'{}\''
             self.project.logger.info(msg.format(self.deps[0].path, self.outs[0].rel_path))
 
@@ -280,7 +298,10 @@ class Stage(object):
 
     def status(self):
         ret = {}
-        ret.update(self._status(self.deps, 'deps'))
+
+        if not self.locked:
+            ret.update(self._status(self.deps, 'deps'))
+
         ret.update(self._status(self.outs, 'outs'))
 
         if ret or self.changed_md5() or self.is_callback:

tests/test_repro.py

@@ -14,7 +14,7 @@ from google.cloud import storage as gc
 
 from dvc.main import main
 from dvc.command.repro import CmdRepro
-from dvc.project import Project, ReproductionError
+from dvc.project import Project, ReproductionError, NotDvcFileError
 from dvc.utils import file_md5
 from dvc.remote.local import RemoteLOCAL
 from dvc.stage import Stage
@@ -130,6 +130,37 @@ class TestReproChangedDeepData(TestReproChangedData):
         self.assertEqual(len(stages), 3)
 
 
+class TestReproLocked(TestReproChangedData):
+    def test(self):
+        file2 = 'file2'
+        file2_stage = file2 + '.dvc'
+        self.dvc.run(fname=file2_stage,
+                     outs=[file2],
+                     deps=[self.file1, self.CODE],
+                     cmd='python {} {} {}'.format(self.CODE, self.file1, file2))
+
+        self.swap_foo_with_bar()
+
+        ret = main(['lock', file2_stage])
+        self.assertEqual(ret, 0)
+        stages = self.dvc.reproduce(file2_stage)
+        self.assertEqual(len(stages), 0)
+
+        ret = main(['unlock', file2_stage])
+        self.assertEqual(ret, 0)
+        stages = self.dvc.reproduce(file2_stage)
+        self.assertTrue(filecmp.cmp(self.file1, self.BAR))
+        self.assertTrue(filecmp.cmp(file2, self.BAR))
+        self.assertEqual(len(stages), 3)
+
+    def test_non_existing(self):
+        with self.assertRaises(NotDvcFileError):
+            self.dvc.lock_stage('non-existing-stage')
+
+        ret = main(['lock', 'non-existing-stage'])
+        self.assertNotEqual(ret, 0)
+
+
 class TestReproPhony(TestReproChangedData):
     def test(self):
         stage = self.dvc.run(deps=[self.file1])