From 52c08b29263bc47e7dc8a4e8cf2db2313df6640b Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Tue, 24 Apr 2018 13:45:30 +0300
Subject: [PATCH] dvc: checkout single file

Fixes #666

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/cli.py              |  4 ++++
 dvc/command/checkout.py | 10 +++++-----
 dvc/project.py          | 12 ++++++++++--
 tests/test_checkout.py  |  9 ++++++++-
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/dvc/cli.py b/dvc/cli.py
index cf1cd99c0..0ea61e3a5 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -91,6 +91,10 @@ def parse_args(argv=None):
                         'checkout',
                         parents=[parent_parser],
                         help='Checkout data files from cache')
+    checkout_parser.add_argument(
+                        'targets',
+                        nargs='*',
+                        help='DVC files')
     checkout_parser.set_defaults(func=CmdCheckout)
 
     # Run
diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py
index b29b77eda..32f77c178 100644
--- a/dvc/command/checkout.py
+++ b/dvc/command/checkout.py
@@ -1,12 +1,12 @@
 import os
-import stat
-
-from dvc.system import System
 from dvc.command.common.base import CmdBase
-from dvc.logger import Logger
 
 
 class CmdCheckout(CmdBase):
     def run(self):
-        self.project.checkout()
+        if not self.args.targets:
+            self.project.checkout()
+        else:
+            for target in self.args.targets:
+                self.project.checkout(target=target)
         return 0
diff --git a/dvc/project.py b/dvc/project.py
index 070a8c370..de3acae58 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -185,9 +185,17 @@ class Project(object):
                 Logger.info(u'Remove empty directory \'{}\''.format(dir))
                 os.removedirs(dir)
 
-    def checkout(self):
+    def checkout(self, target=None):
         self._remove_untracked_hardlinks()
-        for stage in self.stages():
+
+        if target:
+            if not Stage.is_stage_file(target):
+                raise StageNotFoundError(target)
+            stages = [Stage.load(self, target)]
+        else:
+            stages = self.stages()
+
+        for stage in stages:
             stage.checkout()
 
     def _used_cache(self, target=None):
diff --git a/tests/test_checkout.py b/tests/test_checkout.py
index 784569a5a..8732daf66 100644
--- a/tests/test_checkout.py
+++ b/tests/test_checkout.py
@@ -16,7 +16,7 @@ class TestCheckout(TestRepro):
     def setUp(self):
         super(TestCheckout, self).setUp()
 
-        self.dvc.add(self.DATA_DIR)
+        self.data_dir_stage = self.dvc.add(self.DATA_DIR)
 
         self.orig = 'orig'
         shutil.copy(self.FOO, self.orig)
@@ -35,6 +35,13 @@ class TestCheckout(TestRepro):
         self.assertTrue(filecmp.cmp(self.FOO, self.orig))
 
 
+class TestCheckoutSingleStage(TestCheckout):
+    def test(self):
+        self.dvc.checkout(self.foo_stage.path)
+        self.dvc.checkout(self.data_dir_stage.path)
+        self._test_checkout()
+
+
 class TestCheckoutCorruptedCacheFile(TestRepro):
     def test(self):
         cache = self.foo_stage.outs[0].cache