From ace92cb422ed699e4b123d3eea51dbb09113de1f Mon Sep 17 00:00:00 2001
From: pared <pawelredzynski@gmail.com>
Date: Wed, 27 Mar 2019 20:31:26 +0100
Subject: [PATCH] repo: fix stages gathering for checkout

---
 dvc/command/checkout.py |  8 ++++++++
 dvc/exceptions.py       |  7 +++++++
 dvc/repo/__init__.py    |  9 ++++++++-
 dvc/repo/checkout.py    |  6 +++---
 tests/test_checkout.py  | 27 ++++++++++++++++++++++++++-
 5 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py
index be4710870..0f0e0d5ef 100644
--- a/dvc/command/checkout.py
+++ b/dvc/command/checkout.py
@@ -13,6 +13,7 @@ class CmdCheckout(CmdBase):
                     target=target,
                     with_deps=self.args.with_deps,
                     force=self.args.force,
+                    recursive=self.args.recursive,
                 )
         return 0
 
@@ -42,5 +43,12 @@ def add_parser(subparsers, parent_parser):
         default=False,
         help="Do not prompt when removing working directory files.",
     )
+    checkout_parser.add_argument(
+        "-R",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Checkout all subdirectories of the specified directory.",
+    )
     checkout_parser.add_argument("targets", nargs="*", help="DVC files.")
     checkout_parser.set_defaults(func=CmdCheckout)
diff --git a/dvc/exceptions.py b/dvc/exceptions.py
index 0a1df2a6e..57410a892 100644
--- a/dvc/exceptions.py
+++ b/dvc/exceptions.py
@@ -241,3 +241,10 @@ class OverlappingOutputPathsError(DvcException):
                 out_2.stage.relpath,
             )
         )
+
+
+class TargetNotDirectoryError(DvcException):
+    def __init__(self, path):
+        super(TargetNotDirectoryError, self).__init__(
+            "Target: {} is not a directory".format(path)
+        )
diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 5fe31d1aa..69fecde6f 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -5,7 +5,12 @@ import os
 import dvc.prompt as prompt
 import dvc.logger as logger
 
-from dvc.exceptions import DvcException, NotDvcRepoError, OutputNotFoundError
+from dvc.exceptions import (
+    DvcException,
+    NotDvcRepoError,
+    OutputNotFoundError,
+    TargetNotDirectoryError,
+)
 
 
 class Repo(object):
@@ -367,6 +372,8 @@ class Repo(object):
 
         if not from_directory:
             from_directory = self.root_dir
+        elif not os.path.isdir(from_directory):
+            raise TargetNotDirectoryError(from_directory)
 
         stages = []
         outs = []
diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py
index 21a350634..479e05cbc 100644
--- a/dvc/repo/checkout.py
+++ b/dvc/repo/checkout.py
@@ -27,8 +27,9 @@ def get_progress_callback(stages):
 def checkout(self, target=None, with_deps=False, force=False, recursive=False):
     from dvc.stage import StageFileDoesNotExistError, StageFileBadNameError
 
+    all_stages = self.stages()
+
     if target and not recursive:
-        all_stages = self.active_stages()
         try:
             stages = self.collect(target, with_deps=with_deps)
         except (StageFileDoesNotExistError, StageFileBadNameError) as exc:
@@ -36,8 +37,7 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False):
                 str(exc) + " Did you mean 'git checkout {}'?".format(target)
             )
     else:
-        all_stages = self.active_stages(target)
-        stages = all_stages
+        stages = self.active_stages(target)
 
     with self.state:
         _cleanup_unused_links(self, all_stages)
diff --git a/tests/test_checkout.py b/tests/test_checkout.py
index 0055d0941..a07d24bc6 100644
--- a/tests/test_checkout.py
+++ b/tests/test_checkout.py
@@ -17,7 +17,11 @@ from tests.basic_env import TestDvc
 from tests.test_repro import TestRepro
 from dvc.stage import Stage
 from dvc.remote.local import RemoteLOCAL
-from dvc.exceptions import DvcException, ConfirmRemoveError
+from dvc.exceptions import (
+    DvcException,
+    ConfirmRemoveError,
+    TargetNotDirectoryError,
+)
 
 from mock import patch
 from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
@@ -510,3 +514,24 @@ class TestCheckoutShouldHaveSelfClearingProgressBar(TestDvc):
 
     def assertAnyEndsWith(self, update_bars, name):
         self.assertTrue(any(ub for ub in update_bars if ub.endswith(name)))
+
+
+class TestCheckoutTargetRecursiveShouldNotRemoveOtherUsedFiles(TestDvc):
+    def test(self):
+        ret = main(["add", self.DATA_DIR, self.FOO, self.BAR])
+        self.assertEqual(0, ret)
+
+        ret = main(["checkout", "-R", self.DATA_DIR])
+        self.assertEqual(0, ret)
+
+        self.assertTrue(os.path.exists(self.FOO))
+        self.assertTrue(os.path.exists(self.BAR))
+
+
+class TestCheckoutRecursiveNotDirectory(TestDvc):
+    def test(self):
+        ret = main(["add", self.FOO])
+        self.assertEqual(0, ret)
+
+        with self.assertRaises(TargetNotDirectoryError):
+            self.dvc.checkout(target=self.FOO, recursive=True)