From d232a15e8ef467eaa6baf5cb4cb37d20297859d8 Mon Sep 17 00:00:00 2001
From: Dmitry <dmitry.petrov@nevesomo.com>
Date: Sat, 6 Jan 2018 04:15:23 -0800
Subject: [PATCH] Refactoring: move find_cache() to Cache class

---
 dvc/cache.py            | 11 +++++++++++
 dvc/command/checkout.py | 15 ++-------------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/dvc/cache.py b/dvc/cache.py
index dbb8417e9..2a5cdea12 100644
--- a/dvc/cache.py
+++ b/dvc/cache.py
@@ -1,5 +1,7 @@
 import os
 
+from dvc.system import System
+
 
 class Cache(object):
     CACHE_DIR = 'cache'
@@ -23,3 +25,12 @@ class Cache(object):
 
     def get(self, md5):
         return os.path.join(self.cache_dir, md5)
+
+    def find_cache(self, files):
+        file_set = set(files)
+        cached = {}
+        for cache_file in self.all():
+            cached_files = list(filter(lambda f: System.samefile(cache_file, f, True), file_set))
+            cached.update(dict((f, os.path.basename(cache_file)) for f in cached_files))
+            file_set = file_set - set(cached_files)
+        return cached
diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py
index c6e5877fc..c8c171d38 100644
--- a/dvc/command/checkout.py
+++ b/dvc/command/checkout.py
@@ -2,7 +2,6 @@ import os
 
 from dvc.command.common.base import CmdBase
 from dvc.logger import Logger
-from dvc.system import System
 
 
 class CmdCheckout(CmdBase):
@@ -12,9 +11,9 @@ class CmdCheckout(CmdBase):
         return 0
 
     def remove_untracked_hardlinks(self):
-        untracked_cache_files = self.untracked_hardlinks_files()
+        untracked = self.project.scm.untracked_files()
 
-        for file in untracked_cache_files:
+        for file in self.project.cache.find_cache(untracked).keys():
             Logger.info(u'Remove \'{}\''.format(file))
             os.remove(file)
 
@@ -23,13 +22,3 @@ class CmdCheckout(CmdBase):
                 Logger.info(u'Remove empty directory \'{}\''.format(dir))
                 os.removedirs(dir)
         pass
-
-    def untracked_hardlinks_files(self):
-        untracked_files = set(self.project.scm.untracked_files())
-        cached_files = []
-        for cache_file in self.project.cache.all():
-            hardlinks = list(filter(lambda f: System.samefile(cache_file, f), untracked_files))
-            cached_files.extend(hardlinks)
-            untracked_files = untracked_files - set(hardlinks)
-
-        return cached_files