From 7720c3554410fc78fb2cf7eacdaf895fe56160e2 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Tue, 6 Nov 2018 12:41:37 +0200
Subject: [PATCH] checkout: add --with-deps support

Fixes #1296

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/cli.py              |  7 +++++++
 dvc/command/checkout.py |  3 ++-
 dvc/project.py          |  7 ++++---
 tests/test_checkout.py  | 15 +++++++++++++++
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/dvc/cli.py b/dvc/cli.py
index 5dd6d7371..c57a103ca 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -184,6 +184,13 @@ def parse_args(argv=None):
                         parents=[parent_parser],
                         description=CHECKOUT_HELP,
                         help=CHECKOUT_HELP)
+    checkout_parser.add_argument(
+                        '-d',
+                        '--with-deps',
+                        action='store_true',
+                        default=False,
+                        help='Checkout all dependencies of the '
+                             'specified target.')
     checkout_parser.add_argument(
                         'targets',
                         nargs='*',
diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py
index 54ce65932..097dfc7e2 100644
--- a/dvc/command/checkout.py
+++ b/dvc/command/checkout.py
@@ -7,5 +7,6 @@ class CmdCheckout(CmdBase):
             self.project.checkout()
         else:
             for target in self.args.targets:
-                self.project.checkout(target=target)
+                self.project.checkout(target=target,
+                                      with_deps=self.args.with_deps)
         return 0
diff --git a/dvc/project.py b/dvc/project.py
index 8682e62c0..0828dbce4 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -449,11 +449,12 @@ class Project(object):
                 used.append(out.path)
         self.state.remove_unused_links(used)
 
-    def checkout(self, target=None):
+    def checkout(self, target=None, with_deps=False):
         all_stages = self.active_stages()
 
         if target:
-            stages = [Stage.load(self, target)]
+            stages = self._collect(target,
+                                   with_deps=with_deps)
         else:
             stages = all_stages
 
@@ -701,7 +702,7 @@ class Project(object):
                    all_tags=all_tags,
                    show_checksums=show_checksums,
                    with_deps=with_deps)
-        self.checkout(target=target)
+        self.checkout(target=target, with_deps=with_deps)
 
     def _local_status(self, target=None):
         status = {}
diff --git a/tests/test_checkout.py b/tests/test_checkout.py
index 776f85f54..b073c0a64 100644
--- a/tests/test_checkout.py
+++ b/tests/test_checkout.py
@@ -246,3 +246,18 @@ class TestCheckoutNotCachedFile(TestDvc):
                              outs_no_cache=['out'])
 
         self.dvc.checkout()
+
+
+class TestCheckoutWithDeps(TestRepro):
+    def test(self):
+        os.unlink(self.FOO)
+        os.unlink(self.file1)
+
+        self.assertFalse(os.path.exists(self.FOO))
+        self.assertFalse(os.path.exists(self.file1))
+
+        ret = main(['checkout', self.file1_stage, '--with-deps'])
+        self.assertEqual(ret, 0)
+
+        self.assertTrue(os.path.exists(self.FOO))
+        self.assertTrue(os.path.exists(self.file1))