Use `git status --porcelain --branch` to improve performance

Based off of changes by @wkentaro to zsh-git-prompt
see: https://github.com/olivierverdier/zsh-git-prompt/pull/65

gitstatus.py

@@ -32,93 +32,109 @@ except SyntaxError:
           w(str(a))
       w(kwd.get("end", "\n"))
 
-
 # change those symbols to whatever you prefer
 symbols = {'ahead of': '↑·', 'behind': '↓·', 'prehash':':'}
 
-from subprocess import Popen, PIPE
-
 import sys
-gitsym = Popen(['git', 'symbolic-ref', 'HEAD'], stdout=PIPE, stderr=PIPE)
-branch, error = gitsym.communicate()
-
-error_string = error.decode('utf-8')
-
-if 'fatal: Not a git repository' in error_string:
-	sys.exit(0)
-
-branch = branch.decode('utf-8').strip()[11:]
-
-res, err = Popen(['git','diff','--name-status'], stdout=PIPE, stderr=PIPE).communicate()
-err_string = err.decode('utf-8')
-
-if 'fatal' in err_string:
-	sys.exit(0)
-
-changed_files = [namestat[0] for namestat in res.splitlines()]
-staged_files = [namestat[0] for namestat in Popen(['git','diff', '--staged','--name-status'], stdout=PIPE).communicate()[0].splitlines()]
-nb_changed = len(changed_files) - changed_files.count('U')
-nb_U = staged_files.count('U')
-nb_staged = len(staged_files) - nb_U
-staged = str(nb_staged)
-conflicts = str(nb_U)
-changed = str(nb_changed)
-status_lines = Popen(['git','status','-s','-uall'],stdout=PIPE).communicate()[0].splitlines()
-untracked_lines = [a for a in map(lambda s: s.decode('utf-8'), status_lines) if a.startswith("??")]
-nb_untracked = len(untracked_lines)
-untracked = str(nb_untracked)
-stashes = Popen(['git','stash','list'],stdout=PIPE).communicate()[0].splitlines()
-nb_stashed = len(stashes)
-stashed = str(nb_stashed)
-
-if not nb_changed and not nb_staged and not nb_U and not nb_untracked and not nb_stashed:
-	clean = '1'
-else:
-	clean = '0'
-
+import re
+import shlex
+from subprocess import Popen, PIPE, check_output
+
+
+def get_tagname_or_hash():
+    """return tagname if exists else hash"""
+    cmd = 'git log -1 --format="%h%d"'
+    output = check_output(shlex.split(cmd)).decode('utf-8').strip()
+    hash_, tagname = None, None
+    # get hash
+    m = re.search('\(.*\)$', output)
+    if m:
+        hash_ = output[:m.start()-1]
+    # get tagname
+    m = re.search('tag: .*[,\)]', output)
+    if m:
+        tagname = 'tags/' + output[m.start()+len('tag: '): m.end()-1]
+
+    if tagname:
+        return tagname
+    elif hash_:
+        return hash_
+    return None
+
+def get_stash():
+    cmd = Popen(['git', 'rev-parse', '--git-dir'], stdout=PIPE, stderr=PIPE)
+    so, se = cmd.communicate()
+    stashFile = '%s%s' % (so.decode('utf-8').rstrip(),'/logs/refs/stash')
+
+    try:
+        with open(stashFile) as f:
+            return sum(1 for _ in f)
+    except IOError:
+        return 0
+
+# `git status --porcelain --branch` can collect all information
+# branch, remote_branch, untracked, staged, changed, conflicts, ahead, behind
+po = Popen(['git', 'status', '--porcelain', '--branch'], stdout=PIPE, stderr=PIPE)
+stdout, sterr = po.communicate()
+if po.returncode != 0:
+    sys.exit(0)  # Not a git repository
+
+# collect git status information
+untracked, staged, changed, conflicts = [], [], [], []
+ahead, behind = 0, 0
 remote = ''
-
-tag, tag_error = Popen(['git', 'describe', '--exact-match'], stdout=PIPE, stderr=PIPE).communicate()
-
-if not branch: # not on any branch
-	if tag: # if we are on a tag, print the tag's name
-		branch = tag
-	else:
-		branch = symbols['prehash']+ Popen(['git','rev-parse','--short','HEAD'], stdout=PIPE).communicate()[0].decode('utf-8')[:-1]
+status = [(line[0], line[1], line[2:]) for line in stdout.decode('utf-8').splitlines()]
+for st in status:
+    if st[0] == '#' and st[1] == '#':
+        if re.search('Initial commit on', st[2]):
+            branch = st[2].split(' ')[-1]
+        elif re.search('no branch', st[2]):  # detached status
+            branch = get_tagname_or_hash()
+        elif 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())
+                        remote += '%s%s' % (symbols['ahead of'], ahead)
+                    elif 'behind' in div:
+                        behind = int(div[len('behind '):].strip())
+                        remote += '%s%s' % (symbols['behind'], behind)
+    elif st[0] == '?' and st[1] == '?':
+        untracked.append(st)
+    else:
+        if st[1] == 'M':
+            changed.append(st)
+        if st[0] == 'U':
+            conflicts.append(st)
+        elif st[0] != ' ':
+            staged.append(st)
+
+if not changed and not staged and not conflicts and not untracked:
+    clean = 1
 else:
-	remote_name = Popen(['git','config','branch.%s.remote' % branch], stdout=PIPE).communicate()[0].strip()
-	if remote_name:
-		merge_name = Popen(['git','config','branch.%s.merge' % branch], stdout=PIPE).communicate()[0].strip()
-	else:
-		remote_name = "origin"
-		merge_name = "refs/heads/%s" % branch
-
-	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.splitlines()
-	ahead = len([x for x in behead if x[0]=='>'])
-	behind = len(behead) - ahead
-	if behind:
-		remote += '%s%s' % (symbols['behind'], behind)
-	if ahead:
-		remote += '%s%s' % (symbols['ahead of'], ahead)
+    clean = 0
 
 if remote == "":
-	remote = '.'
+    remote = '.'
 
 out = '\n'.join([
-	str(branch),
-	str(remote),
-	staged,
-	conflicts,
-	changed,
-	untracked,
-	stashed,
-	clean])
+    branch,
+    remote.decode('utf-8'),
+    str(len(staged)),
+    str(len(conflicts)),
+    str(len(changed)),
+    str(len(untracked)),
+    str(get_stash()),
+    str(clean)
+])
 Print(out)

gitstatus.sh

@@ -6,8 +6,7 @@
 # Alan K. Stebbens <aks@stebbens.org> [http://github.com/aks]
 
 # helper functions
-count_lines() { echo "$1" | egrep -c "^$2" ; }
-all_lines() { echo "$1" | grep -v "^$" | wc -l ; }
+count_lines() { echo "$1" | egrep -c $3 "^$2" ; }
 
 if [ -z "${__GIT_PROMPT_DIR}" ]; then
   SOURCE="${BASH_SOURCE[0]}"
@@ -19,39 +18,38 @@ if [ -z "${__GIT_PROMPT_DIR}" ]; then
   __GIT_PROMPT_DIR="$( cd -P "$( dirname "${SOURCE}" )" && pwd )"
 fi
 
-gitsym=`git symbolic-ref HEAD`
+gitstatus=`git status --porcelain --branch`
 
-# if "fatal: Not a git repo .., then exit
-case "$gitsym" in fatal*) exit 0 ;; esac
+# if the status is fatal, exit now
+[[ "$?" -ne 0 ]] && exit 0
 
-# the current branch is the tail end of the symbolic reference
-branch="${gitsym##refs/heads/}"    # get the basename after "refs/heads/"
+num_staged=`count_lines "$gitstatus" "(\?\?|##| )" "-v"`
+num_changed=`count_lines "$gitstatus" ".M"`
+num_conflicts=`count_lines "$gitstatus" "U"`
+num_untracked=`count_lines "$gitstatus" "\?\?"`
 
-gitstatus=`git diff --name-status 2>&1`
-
-# if the diff is fatal, exit now
-case "$gitstatus" in fatal*) exit 0 ;; esac
-
-
-staged_files=`git diff --staged --name-status`
-
-num_changed=$(( `all_lines "$gitstatus"` - `count_lines "$gitstatus" U` ))
-num_conflicts=`count_lines "$staged_files" U`
-num_staged=$(( `all_lines "$staged_files"` - num_conflicts ))
-num_untracked=`git ls-files --others --exclude-standard $(git rev-parse --show-cdup) | wc -l`
 if [[ "$__GIT_PROMPT_IGNORE_STASH" = "1" ]]; then
   num_stashed=0
-else	
-  num_stashed=`git stash list | wc -l`
+else
+  stash_file="`git rev-parse --git-dir`/logs/refs/stash"
+  if [[ -e "${stash_file}" ]]; then
+    num_stashed=`wc -l "${stash_file}" | cut -d' ' -f 1`
+  else
+    num_stashed=0
+  fi
 fi
 
 clean=0
-if (( num_changed == 0 && num_staged == 0 && num_U == 0 && num_untracked == 0 && num_stashed == 0 )) ; then
+if (( num_changed == 0 && num_staged == 0 && num_untracked == 0 && num_stashed == 0 )) ; then
   clean=1
 fi
 
 remote=
 
+branch_line=`echo $gitstatus | grep "^##"`
+IFS="." read -ra line <<< "${branch_line/\#\# }"
+branch="${line[0]}"
+
 if [[ -z "$branch" ]]; then
   tag=`git describe --exact-match`
   if [[ -n "$tag" ]]; then
@@ -59,52 +57,27 @@ if [[ -z "$branch" ]]; then
   else
     branch="_PREHASH_`git rev-parse --short HEAD`"
   fi
+elif [[ "$branch" == *"Initial commit on"* ]]; then
+  IFS=" " read -ra branch_line <<< "$branch"
+  branch=${branch_line[-1]}
+elif [[ "$branch" == *"no branch"* ]]; then
+  branch="_PREHASH_`git rev-parse --short HEAD`"
 else
-  remote_name=`git config branch.${branch}.remote`
-
-  if [[ -n "$remote_name" ]]; then
-    merge_name=`git config branch.${branch}.merge`
-  else
-    remote_name='origin'
-    merge_name="refs/heads/${branch}"
-  fi
-
-  if [[ "$remote_name" == '.' ]]; then
-    remote_ref="$merge_name"
-  else
-    remote_ref="refs/remotes/$remote_name/${merge_name##refs/heads/}"
-  fi
-
-  # detect if the local branch have a remote tracking branch
-  cmd_output=$(git rev-parse --abbrev-ref ${branch}@{upstream} 2>&1 >/dev/null)
-
-  if [ `count_lines "$cmd_output" "fatal: no upstream"` == 1 ] ; then
-    has_remote_tracking=0
-  else
-    has_remote_tracking=1
+  IFS="[]" read -ra remote_line <<< "${line[3]}"
+  if [[ "${remote_line[1]}" == *ahead* ]]; then
+    num_ahead=`echo "${remote_line[1]}" | cut -c 7-`
+    remote="${remote}_AHEAD_${num_ahead}"
   fi
-
-  # get the revision list, and count the leading "<" and ">"
-  revgit=`git rev-list --left-right ${remote_ref}...HEAD`
-  num_revs=`all_lines "$revgit"`
-  num_ahead=`count_lines "$revgit" "^>"`
-  num_behind=$(( num_revs - num_ahead ))
-  if (( num_behind > 0 )) ; then
+  if [[ "${remote_line[1]}" == *behind* ]]; then
+    num_behind=`echo "${remote_line[1]}" | cut -c 9-`
     remote="${remote}_BEHIND_${num_behind}"
   fi
-  if (( num_ahead > 0 )) ; then
-    remote="${remote}_AHEAD_${num_ahead}"
-  fi
 fi
 
 if [[ -z "$remote" ]] ; then
   remote='.'
 fi
 
-if [[ "$has_remote_tracking" == "0" ]] ; then
-  remote='_NO_REMOTE_TRACKING_'
-fi 
-
 for w in "$branch" "$remote" $num_staged $num_conflicts $num_changed $num_untracked $num_stashed $clean ; do
   echo "$w"
 done