Merge pull request #288 from dataversioncontrol/dvc_checkout_fix

dvc checkout fix

dvc/command/checkout.py

@@ -4,6 +4,7 @@ from dvc.command.base import CmdBase, DvcLock
 from dvc.logger import Logger
 from dvc.system import System
 
+
 class CmdCheckout(CmdBase):
     def __init__(self, settings):
         super(CmdCheckout, self).__init__(settings)
@@ -34,6 +35,30 @@ class CmdCheckout(CmdBase):
  
     def run(self):
         with DvcLock(self.is_locker, self.git):
-            items = self.settings.path_factory.all_existing_data_items()
-            self.checkout(items)
+            curr_commit = self.git.curr_branch_or_commit
+            prev_items = []
+            try:
+                self.git.checkout_previous()
+                prev_items = self.settings.path_factory.all_existing_data_items()
+            except Exception as ex:
+                Logger.error(u'Unable to get '.format(ex))
+                return 1
+            finally:
+                self.git.checkout(curr_commit)
+
+            curr_items = self.settings.path_factory.all_existing_data_items()
+            self.checkout(curr_items)
+
+            self.remove_files(list(set(prev_items) - set(curr_items)))
             return 0
+
+    @staticmethod
+    def remove_files(removed_items_set):
+        for item in removed_items_set:
+            Logger.info(u'Remove \'{}\''.format(item.data.relative))
+            os.remove(item.data.relative)
+            dir = os.path.dirname(item.data.relative)
+            if not os.listdir(dir):
+                Logger.info(u'Remove directory \'{}\''.format(dir))
+                os.removedirs(dir)
+        pass

dvc/git_wrapper.py

@@ -189,16 +189,33 @@ class GitWrapper(GitWrapperI):
             raise ExecutorError('Unable to run git command: {}'.format(e))
         pass
 
-    @property
-    def curr_commit(self):
+    def curr_commit(self, branches=False):
         Logger.debug('[dvc-git] Getting current git commit. Command: git rev-parse --short HEAD')
 
-        code, out, err = Executor.exec_cmd(['git', 'rev-parse', 'HEAD'])
+        if branches:
+            cmd = 'git rev-parse --abbrev-ref HEAD'
+        else:
+            cmd = 'git rev-parse HEAD'
+
+        code, out, err = Executor.exec_cmd(cmd.split())
         if code != 0:
             raise ExecutorError('[dvc-git] Commit command error - {}'.format(err))
         Logger.debug('[dvc-git] Getting current git commit. Success.')
         return out
 
+    @property
+    def curr_branch_or_commit(self):
+        Logger.debug('[dvc-git] Getting current git branch or commit.')
+
+        branch = self.curr_commit(branches=True)
+        if branch == '' or branch == 'HEAD':
+            result = self.curr_commit(branches=False)
+        else:
+            result = branch
+
+        Logger.debug('[dvc-git] Getting current git branch or commit. Result {}. Success.'.format(result))
+        return result
+
     @staticmethod
     def commit_all_changes(message):
         Logger.debug('[dvc-git] Commit all changes. Commands: {} && {} && {}'.format(
@@ -217,7 +234,7 @@ class GitWrapper(GitWrapperI):
     def commit_all_changes_and_log_status(self, message):
         statuses = self.commit_all_changes(message)
         Logger.debug('[dvc-git] A new commit {} was made in the current branch. '
-                     'Added files:'.format(self.curr_commit))
+                     'Added files:'.format(self.curr_commit()))
         for status, file in statuses:
             Logger.debug('[dvc-git]\t{} {}'.format(status, file))
         pass
@@ -299,6 +316,18 @@ class GitWrapper(GitWrapperI):
         Logger.debug('[dvc-git] Checkout file {} before last merge. Command: {}'.format(fname, cmd))
         Executor.exec_cmd_only_success(cmd.split())
 
+    @staticmethod
+    def checkout(commit):
+        Logger.debug('[dvc-git] Checkout {}'.format(commit))
+        cmd = 'git checkout {}'.format(commit)
+        Executor.exec_cmd(cmd.split())
+
+    @staticmethod
+    def checkout_previous():
+        Logger.debug('[dvc-git] Checkout previous')
+        cmd = 'git checkout -'
+        Executor.exec_cmd(cmd.split())
+
     def separate_dependency_files_and_dirs(self, code_dependencies):
         code_files = []
         code_dirs = []