From b2f0a7654642c4e65b4520ed418476678ffec14f Mon Sep 17 00:00:00 2001
From: Dmitry <dmitry.petrov@nevesomo.com>
Date: Sat, 4 Nov 2017 17:14:44 -0700
Subject: [PATCH] Checkout branches, not commits if possible

---
 dvc/command/checkout.py |  6 +++---
 dvc/git_wrapper.py      | 25 +++++++++++++++++++++----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/dvc/command/checkout.py b/dvc/command/checkout.py
index bdc2dd93b..5e9fa3ab9 100644
--- a/dvc/command/checkout.py
+++ b/dvc/command/checkout.py
@@ -35,7 +35,7 @@ class CmdCheckout(CmdBase):
  
     def run(self):
         with DvcLock(self.is_locker, self.git):
-            curr_commit = self.git.curr_commit
+            curr_commit = self.git.curr_branch_or_commit
             prev_items = []
             try:
                 self.git.checkout_previous()
@@ -55,10 +55,10 @@ class CmdCheckout(CmdBase):
     @staticmethod
     def remove_files(removed_items_set):
         for item in removed_items_set:
-            Logger.info(u'Removing data file not from the current branch {}'.format(item.data.relative))
+            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'Removing directory not from the current branch {}'.format(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 391805ab4..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