Restyle commands/update: add --rev option

dvc/command/imp.py

@@ -53,6 +53,6 @@ def add_parser(subparsers, parent_parser):
         "-o", "--out", nargs="?", help="Destination path to download files to"
     )
     import_parser.add_argument(
-        "--rev", nargs="?", help="Git revision (e.g. branch, tag, SHA)"
+        "--rev", nargs="?", help="Git revision in repository to update from."
     )
     import_parser.set_defaults(func=CmdImport)

dvc/command/update.py

@@ -14,7 +14,7 @@ def run(self):
         ret = 0
         for target in self.args.targets:
             try:
-                self.repo.update(target)
+                self.repo.update(target, rev=self.args.rev)
             except DvcException:
                 logger.exception("failed to update '{}'.".format(target))
                 ret = 1
@@ -33,4 +33,7 @@ def add_parser(subparsers, parent_parser):
     update_parser.add_argument(
         "targets", nargs="+", help="DVC-files to update."
     )
+    update_parser.add_argument(
+        "--rev", nargs="?", help="Git revision in repository to update from."
+    )
     update_parser.set_defaults(func=CmdUpdate)

dvc/dependency/base.py

@@ -1,4 +1,5 @@
 from dvc.exceptions import DvcException
+from dvc.exceptions import UpdateWithRevNotPossibleError
 
 
 class DependencyDoesNotExistError(DvcException):
@@ -27,5 +28,6 @@ class DependencyBase(object):
     IsNotFileOrDirError = DependencyIsNotFileOrDirError
     IsStageFileError = DependencyIsStageFileError
 
-    def update(self):
-        pass
+    def update(self, rev=None):
+        if rev:
+            raise UpdateWithRevNotPossibleError()

dvc/dependency/repo.py

@@ -122,6 +122,9 @@ def download(self, to):
                 self.def_path, self.def_repo[self.PARAM_URL]
             )
 
-    def update(self):
-        with self._make_repo(rev_lock=None) as repo:
+    def update(self, rev=None):
+        if not rev:
+            rev = self.def_repo.get("rev")
+
+        with self._make_repo(rev=rev, rev_lock=None) as repo:
             self.def_repo[self.PARAM_REV_LOCK] = repo.scm.get_rev()

dvc/exceptions.py

@@ -309,3 +309,10 @@ def __init__(self, path, repo):
             " neighther as an output nor a git-handled file."
         )
         super().__init__(msg.format(path, repo))
+
+
+class UpdateWithRevNotPossibleError(DvcException):
+    def __init__(self):
+        super(UpdateWithRevNotPossibleError, self).__init__(
+            "Revision option (`--rev`) is not a feature of non-Git sources."
+        )

dvc/repo/update.py

@@ -2,10 +2,11 @@
 
 
 @locked
-def update(self, target):
+def update(self, target, rev=None):
     from dvc.stage import Stage
 
     stage = Stage.load(self, target)
-    stage.update()
+
+    stage.update(rev=rev)
 
     stage.dump()

dvc/stage.py

@@ -405,11 +405,11 @@ def reproduce(self, interactive=False, **kwargs):
 
         return self
 
-    def update(self):
+    def update(self, rev=None):
         if not self.is_repo_import and not self.is_import:
             raise StageUpdateError(self.relpath)
 
-        self.deps[0].update()
+        self.deps[0].update(rev)
         locked = self.locked
         self.locked = False
         try:

tests/func/test_update.py

@@ -1,8 +1,13 @@
+import filecmp
+import os
+import shutil
 import pytest
 
-from dvc.stage import Stage
-from dvc.compat import fspath
+from dvc.exceptions import UpdateWithRevNotPossibleError
 from dvc.external_repo import clean_repos
+from dvc.repo import Repo
+from dvc.compat import fspath
+from dvc.stage import Stage
 
 
 @pytest.mark.parametrize("cached", [True, False])
@@ -71,3 +76,33 @@ def test_update_import_url(tmp_dir, dvc, tmp_path_factory):
 
     assert dst.is_file()
     assert dst.read_text() == "updated file content"
+
+
+def test_update_rev(tmp_dir, dvc, erepo_dir, monkeypatch):
+    with monkeypatch.context() as m:
+        m.chdir(fspath(erepo_dir))
+        erepo_dir.scm.checkout("new_branch", create_new=True)
+        erepo_dir.dvc_gen("foo", "foo content", commit="create foo on branch")
+        erepo_dir.scm.checkout("master")
+        erepo_dir.scm.checkout("new_branch_2", create_new=True)
+        erepo_dir.dvc_gen(
+            "foo", "foo content 2", commit="create foo on branch"
+        )
+        erepo_dir.scm.checkout("master")
+
+    stage = dvc.imp(fspath(erepo_dir), "foo", "foo_imported", rev="new_branch")
+    dvc.update(stage.path, rev="new_branch_2")
+
+    assert (tmp_dir / "foo_imported").read_text() == "foo content 2"
+
+
+def test_update_rev_non_git_failure(repo_dir, dvc_repo):
+    src = "file"
+    dst = src + "_imported"
+
+    shutil.copyfile(repo_dir.FOO, src)
+
+    stage = dvc_repo.imp_url(src, dst)
+
+    with pytest.raises(UpdateWithRevNotPossibleError):
+        dvc_repo.update(stage.path, rev="dev")

tests/unit/command/test_update.py

@@ -4,12 +4,12 @@
 
 def test_update(dvc, mocker):
     targets = ["target1", "target2", "target3"]
-    cli_args = parse_args(["update"] + targets)
+    cli_args = parse_args(["update", "--rev", "develop"] + targets)
     assert cli_args.func == CmdUpdate
     cmd = cli_args.func(cli_args)
     m = mocker.patch("dvc.repo.Repo.update")
 
     assert cmd.run() == 0
 
-    calls = [mocker.call(target) for target in targets]
+    calls = [mocker.call(target, rev="develop") for target in targets]
     m.assert_has_calls(calls)