From 94007f7db2d47f3e13e4aac878f0e780a479f628 Mon Sep 17 00:00:00 2001
From: Kentaro Wada <www.kentaro.wada@gmail.com>
Date: Sun, 2 Aug 2015 08:50:56 +0900
Subject: [PATCH] git-prompt: collect all git status information by one cmd

This commit uses the `git status --porcelain -b` format to gather the
number of ahead and behind commits related to the remote branch.
---
 plugins/git-prompt/gitstatus.py | 61 +++++++++++++++------------------
 1 file changed, 28 insertions(+), 33 deletions(-)

diff --git a/plugins/git-prompt/gitstatus.py b/plugins/git-prompt/gitstatus.py
index 903ad474..2f02e5fc 100644
--- a/plugins/git-prompt/gitstatus.py
+++ b/plugins/git-prompt/gitstatus.py
@@ -5,25 +5,40 @@ from __future__ import print_function
 prehash = ':'
 
 import sys
-import subprocess
 from subprocess import Popen, PIPE
 
-branch, error = Popen(['git', 'symbolic-ref', 'HEAD'], stdout=PIPE, stderr=PIPE).communicate()
 
-if 'fatal: Not a git repository' in error.decode('utf-8'):
-    sys.exit(0)
+# `git status --porcelain -b` can collect all information
+# branch, remote_branch, untracked, staged, changed, conflicts, ahead, behind
+po = Popen(['git', 'status', '--porcelain', '-b'], stdout=PIPE, stderr=PIPE)
+stdout, sterr = po.communicate()
+if po.returncode != 0:
+    sys.exit(0)  # Not a git repository
 
-branch = branch.decode("utf-8").strip()[11:]
-
-# Get git status (staged, change, conflicts and untracked)
-try:
-    res = subprocess.check_output(['git', 'status', '--porcelain'])
-except subprocess.CalledProcessError:
-    sys.exit(0)
-status = [(st[0], st[1], st[2:]) for st in res.splitlines()]
+# collect git status information
 untracked, staged, changed, conflicts = [], [], [], []
+ahead, behind = 0, 0
+status = [(line[0], line[1], line[2:]) for line in stdout.splitlines()]
 for st in status:
-    if st[0] == '?' and st[1] == '?':
+    if st[0] == '#' and st[1] == '#':
+        if len(st[2].strip().split('...')) == 1:
+            branch = st[2].strip()
+        else:
+            # current and remote branch info
+            branch, rest = st[2].strip().split('...')
+            if len(rest.split(' ')) == 1:
+                # remote_branch = rest.split(' ')[0]
+                pass
+            else:
+                # ahead or behind
+                divergence = ' '.join(rest.split(' ')[1:])
+                divergence = divergence.lstrip('[').rstrip(']')
+                for div in divergence.split(', '):
+                    if 'ahead' in div:
+                        ahead = int(div[len('ahead '):].strip())
+                    elif 'behind' in div:
+                        behind = int(div[len('behind '):].strip())
+    elif st[0] == '?' and st[1] == '?':
         untracked.append(st)
     else:
         if st[1] == 'M':
@@ -33,26 +48,6 @@ for st in status:
         elif st[0] != ' ':
             staged.append(st)
 
-ahead, behind = 0,0
-
-if not branch: # not on any branch
-    branch = prehash + Popen(['git','rev-parse','--short','HEAD'], stdout=PIPE).communicate()[0].decode("utf-8")[:-1]
-else:
-    remote_name = Popen(['git','config','branch.%s.remote' % branch], stdout=PIPE).communicate()[0].decode("utf-8").strip()
-    if remote_name:
-        merge_name = Popen(['git','config','branch.%s.merge' % branch], stdout=PIPE).communicate()[0].decode("utf-8").strip()
-        if remote_name == '.': # local
-            remote_ref = merge_name
-        else:
-            remote_ref = 'refs/remotes/%s/%s' % (remote_name, merge_name[11:])
-        revgit = Popen(['git', 'rev-list', '--left-right', '%s...HEAD' % remote_ref],stdout=PIPE, stderr=PIPE)
-        revlist = revgit.communicate()[0]
-        if revgit.poll(): # fallback to local
-            revlist = Popen(['git', 'rev-list', '--left-right', '%s...HEAD' % merge_name],stdout=PIPE, stderr=PIPE).communicate()[0]
-        behead = revlist.decode("utf-8").splitlines()
-        ahead = len([x for x in behead if x[0]=='>'])
-        behind = len(behead) - ahead
-
 out = ' '.join([
     branch,
     str(ahead),