dvc: store list of links to cleanup later

Fixes #685

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/cache.py

@@ -2,7 +2,7 @@ import os
 import json
 import shutil
 
-from dvc.state import State
+from dvc.state import State, LinkState
 from dvc.system import System
 from dvc.logger import Logger
 from dvc.utils import move, remove
@@ -21,7 +21,7 @@ class Cache(object):
         'reflink': System.reflink,
     }
 
-    def __init__(self, dvc_dir, cache_dir=None, cache_type=None):
+    def __init__(self, root_dir, dvc_dir, cache_dir=None, cache_type=None):
         self.cache_type = cache_type
 
         cache_dir = cache_dir if cache_dir else self.CACHE_DIR
@@ -34,11 +34,12 @@ class Cache(object):
             os.mkdir(self.cache_dir)
 
         self.state = State(self.cache_dir)
+        self.link_state = LinkState(root_dir, dvc_dir)
         self.lock = Lock(self.cache_dir, name=self.CACHE_DIR_LOCK)
 
     @staticmethod
-    def init(dvc_dir, cache_dir=None):
-        return Cache(dvc_dir, cache_dir=None)
+    def init(root_dir, dvc_dir, cache_dir=None):
+        return Cache(root_dir, dvc_dir, cache_dir=None)
 
     def all(self):
         with self.lock:
@@ -91,6 +92,7 @@ class Cache(object):
         for typ in types:
             try:
                 self.CACHE_TYPE_MAP[typ](src, link)
+                self.link_state.update(link)
                 return
             except Exception as exc:
                 msg = 'Cache type \'{}\' is not supported'.format(typ)

dvc/project.py

@@ -47,7 +47,7 @@ class Project(object):
         self.config = Config(self.dvc_dir)
         self.scm = SCM(self.root_dir)
         self.lock = Lock(self.dvc_dir)
-        self.cache = Cache(self.dvc_dir, cache_dir=self.config._config[Config.SECTION_CACHE].get(Config.SECTION_CACHE_DIR, None),
+        self.cache = Cache(self.root_dir, self.dvc_dir, cache_dir=self.config._config[Config.SECTION_CACHE].get(Config.SECTION_CACHE_DIR, None),
                                          cache_type=self.config._config[Config.SECTION_CACHE].get(Config.SECTION_CACHE_TYPE, None))
         self.state = State(self.dvc_dir)
         self.logger = Logger(self.config._config[Config.SECTION_CORE].get(Config.SECTION_CORE_LOGLEVEL, None))
@@ -81,12 +81,14 @@ class Project(object):
         os.mkdir(dvc_dir)
 
         config = Config.init(dvc_dir)
-        cache = Cache.init(dvc_dir)
+        cache = Cache.init(root_dir, dvc_dir)
         state = State.init(dvc_dir)
         lock = Lock(dvc_dir)
         updater = Updater(dvc_dir)
 
         scm.ignore_list([cache.cache_dir,
+                         cache.link_state.state_file,
+                         cache.link_state._lock_file.lock_file,
                          state.state_file,
                          state._lock_file.lock_file,
                          lock.lock_file,
@@ -177,30 +179,13 @@ class Project(object):
                 raise ReproductionError(stages[n].relpath, ex)
         return result
 
-    def _remove_untracked_hardlinks(self):
-        untracked = self.scm.untracked_files()
-        cache = dict((System.inode(c), c) for c in self.cache.all())
-        for file in untracked:
-            inode = System.inode(file)
-            if inode not in cache.keys():
-                continue
-
-            Logger.info(u'Remove \'{}\''.format(file))
-            os.remove(file)
-
-            dir = os.path.dirname(file)
-            if len(dir) != 0 and not os.listdir(dir):
-                Logger.info(u'Remove empty directory \'{}\''.format(dir))
-                os.removedirs(dir)
-
     def checkout(self, target=None):
-        self._remove_untracked_hardlinks()
-
         if target:
             if not Stage.is_stage_file(target):
                 raise StageNotFoundError(target)
             stages = [Stage.load(self, target)]
         else:
+            self.cache.link_state.remove_all()
             stages = self.stages()
 
         for stage in stages:

dvc/state.py

@@ -5,7 +5,7 @@ import threading
 
 from dvc.lock import Lock
 from dvc.system import System
-from dvc.utils import file_md5, dict_md5
+from dvc.utils import file_md5, dict_md5, remove
 from dvc.exceptions import DvcException
 
 
@@ -149,3 +149,69 @@ class State(object):
         inode = self.inode(path)
 
         return self._get(inode, mtime)
+
+
+class LinkStateEntry(object):
+    PARAM_MTIME = 'mtime'
+    PARAM_INODE = 'inode'
+
+    def __init__(self, inode, mtime):
+        self.mtime = mtime
+        self.inode = inode
+
+    @staticmethod
+    def loadd(d):
+        mtime = d[LinkStateEntry.PARAM_MTIME]
+        inode = d[LinkStateEntry.PARAM_INODE]
+        return LinkStateEntry(inode, mtime)
+
+    def dumpd(self):
+        return {
+            self.PARAM_INODE: self.inode,
+            self.PARAM_MTIME: self.mtime,
+        }
+
+
+class LinkState(State):
+    STATE_FILE = 'link.state'
+    STATE_LOCK_FILE = STATE_FILE + '.lock'
+
+    def __init__(self, root_dir, dvc_dir):
+        super(LinkState, self).__init__(dvc_dir)
+        self.root_dir = root_dir
+
+    def update(self, path):
+        if not os.path.exists(path):
+            return
+
+        mtime = self.mtime(path)
+        inode = self.inode(path)
+
+        with self._lock:
+            with self._lock_file:
+                state = LinkStateEntry(inode, mtime)
+                d = state.dumpd()
+                self._db[os.path.relpath(path, self.root_dir)] = d
+                self.dump()
+
+    def _do_remove_all(self):
+        for p, s in self._db.items():
+            path = os.path.join(self.root_dir, p)
+            state = LinkStateEntry.loadd(s)
+
+            if not os.path.exists(path):
+                continue
+
+            inode = self.inode(path)
+            mtime = self.mtime(path)
+
+            if inode == state.inode and mtime == state.mtime:
+                remove(path)
+
+        self._db = {}
+
+    def remove_all(self):
+        with self._lock:
+            with self._lock_file:
+                self._do_remove_all()
+                self.dump()

tests/test_cache.py

@@ -17,13 +17,13 @@ class TestCache(TestDvc):
         self.create(self.cache2, '2')
 
     def test_all(self):
-        flist = Cache(self.dvc.dvc_dir).all()
+        flist = Cache(self.dvc.root_dir, self.dvc.dvc_dir).all()
         self.assertEquals(len(flist), 2)
         self.assertTrue(self.cache1 in flist)
         self.assertTrue(self.cache2 in flist)
 
     def test_get(self):
-        cache = Cache(self.dvc.dvc_dir).get(self.cache1_md5)
+        cache = Cache(self.dvc.root_dir, self.dvc.dvc_dir).get(self.cache1_md5)
         self.assertEquals(cache, self.cache1)