From 47c04d921ebb3fe4916eb131ae11557650aa1ae3 Mon Sep 17 00:00:00 2001
From: Arnaud Thimel <thimel@codelutin.com>
Date: Thu, 25 Feb 2016 16:36:31 +0100
Subject: [PATCH] feat(git-prompt): display stash count in prompt

---
 plugins/git-prompt/git-prompt.plugin.zsh |  8 +++++++-
 plugins/git-prompt/gitstatus.py          | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/plugins/git-prompt/git-prompt.plugin.zsh b/plugins/git-prompt/git-prompt.plugin.zsh
index 43a327ab0..a4c952471 100644
--- a/plugins/git-prompt/git-prompt.plugin.zsh
+++ b/plugins/git-prompt/git-prompt.plugin.zsh
@@ -40,6 +40,8 @@ function update_current_git_vars() {
     GIT_CONFLICTS=$__CURRENT_GIT_STATUS[5]
     GIT_CHANGED=$__CURRENT_GIT_STATUS[6]
     GIT_UNTRACKED=$__CURRENT_GIT_STATUS[7]
+    GIT_STASHED=$__CURRENT_GIT_STATUS[8]
+    GIT_CLEAN=$__CURRENT_GIT_STATUS[9]
 }
 
 git_super_status() {
@@ -65,7 +67,10 @@ git_super_status() {
       if [ "$GIT_UNTRACKED" -ne "0" ]; then
           STATUS="$STATUS$ZSH_THEME_GIT_PROMPT_UNTRACKED$GIT_UNTRACKED%{${reset_color}%}"
       fi
-      if [ "$GIT_CHANGED" -eq "0" ] && [ "$GIT_CONFLICTS" -eq "0" ] && [ "$GIT_STAGED" -eq "0" ] && [ "$GIT_UNTRACKED" -eq "0" ]; then
+      if [ "$GIT_STASHED" -ne "0" ]; then
+          STATUS="$STATUS$ZSH_THEME_GIT_PROMPT_STASHED$GIT_STASHED%{${reset_color}%}"
+      fi
+      if [ "$GIT_CLEAN" -eq "1" ]; then
           STATUS="$STATUS$ZSH_THEME_GIT_PROMPT_CLEAN"
       fi
       STATUS="$STATUS%{${reset_color}%}$ZSH_THEME_GIT_PROMPT_SUFFIX"
@@ -84,6 +89,7 @@ ZSH_THEME_GIT_PROMPT_CHANGED="%{$fg[blue]%}%{✚%G%}"
 ZSH_THEME_GIT_PROMPT_BEHIND="%{↓%G%}"
 ZSH_THEME_GIT_PROMPT_AHEAD="%{↑%G%}"
 ZSH_THEME_GIT_PROMPT_UNTRACKED="%{$fg[cyan]%}%{…%G%}"
+ZSH_THEME_GIT_PROMPT_STASHED="%{$fg_bold[blue]%}%{⚑%G%}"
 ZSH_THEME_GIT_PROMPT_CLEAN="%{$fg_bold[green]%}%{✔%G%}"
 
 # Set the prompt.
diff --git a/plugins/git-prompt/gitstatus.py b/plugins/git-prompt/gitstatus.py
index bf3173614..786274a71 100644
--- a/plugins/git-prompt/gitstatus.py
+++ b/plugins/git-prompt/gitstatus.py
@@ -23,6 +23,18 @@ def get_tagname_or_hash():
         return hash_
     return None
 
+# Re-use method from https://github.com/magicmonty/bash-git-prompt to get stashs count
+def get_stash():
+    cmd = Popen(['git', 'rev-parse', '--git-dir'], stdout=PIPE, stderr=PIPE)
+    so, se = cmd.communicate()
+    stash_file = '%s%s' % (so.decode('utf-8').rstrip(), '/logs/refs/stash')
+
+    try:
+        with open(stash_file) 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
@@ -68,6 +80,12 @@ for st in status:
         elif st[0] != ' ':
             staged.append(st)
 
+stashed = get_stash()
+if not changed and not staged and not conflicts and not untracked and not stashed:
+    clean = 1
+else:
+    clean = 0
+
 out = ' '.join([
     branch,
     str(ahead),
@@ -76,5 +94,7 @@ out = ' '.join([
     str(len(conflicts)),
     str(len(changed)),
     str(len(untracked)),
+    str(stashed),
+    str(clean)
 ])
 print(out, end='')