diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py index b135642ab..5e9fa3ab9 100644 --- a/dvc/command/checkout.py +++ b/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 diff --git a/dvc/git_wrapper.py b/dvc/git_wrapper.py index e142bc8b6..29e590dc3 100644 --- a/dvc/git_wrapper.py +++ b/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 = []