From 6df97363b82f3bdeee6728a6e722947a0bf2f91f Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Sat, 10 Oct 2020 19:27:04 +0200
Subject: [PATCH] refactor: display git stash count

---
 docs/docs/segment-git.md |  2 ++
 segment_git.go           | 28 ++++++++++------------------
 segment_git_test.go      | 24 ++++--------------------
 3 files changed, 16 insertions(+), 38 deletions(-)

diff --git a/docs/docs/segment-git.md b/docs/docs/segment-git.md
index d0737114..7fc0497b 100644
--- a/docs/docs/segment-git.md
+++ b/docs/docs/segment-git.md
@@ -50,3 +50,5 @@ Local changes can also shown by default using the following syntax for both the
 - cherry_pick_icon: `string` - icon/text to display before the context when doing a cherry-pick
 - commit_icon: `string` - icon/text to display before the commit context
 - tag_icon: `string` - icon/text to display before the tag context
+- display_stash_count: `boolean` show stash count or not
+- stash_count_icon: `string` icon/text to display before the stash context
diff --git a/segment_git.go b/segment_git.go
index 94435a70..6f01dd28 100755
--- a/segment_git.go
+++ b/segment_git.go
@@ -15,7 +15,7 @@ type gitRepo struct {
 	behind     int
 	HEAD       string
 	upstream   string
-	stashCount int
+	stashCount string
 }
 
 type gitStatus struct {
@@ -57,6 +57,10 @@ const (
 	CommitIcon Property = "commit_icon"
 	//TagIcon shows before the tag context
 	TagIcon Property = "tag_icon"
+	//DisplayStashCount show stash count or not
+	DisplayStashCount Property = "display_stash_count"
+	//StashCountIcon shows before the stash context
+	StashCountIcon Property = "stash_count_icon"
 )
 
 func (g *git) enabled() bool {
@@ -97,7 +101,9 @@ func (g *git) string() string {
 	if g.hasWorking() {
 		fmt.Fprintf(buffer, " %s +%d ~%d -%d", g.props.getString(LocalWorkingIcon, "#"), g.repo.working.added+g.repo.working.untracked, g.repo.working.modified, g.repo.working.deleted)
 	}
-	// TODO: Add stash entries
+	if g.props.getBool(DisplayStashCount, false) && g.repo.stashCount != "" {
+		fmt.Fprintf(buffer, " %s%s", g.props.getString(StashCountIcon, ""), g.repo.stashCount)
+	}
 	return buffer.String()
 }
 
@@ -210,9 +216,8 @@ func (g *git) parseGitStats(output []string, working bool) *gitStatus {
 	return &status
 }
 
-func (g *git) getStashContext() int {
-	stash := g.getGitCommandOutput("stash", "list")
-	return numberOfLinesInString(stash)
+func (g *git) getStashContext() string {
+	return g.getGitCommandOutput("rev-list", "--walk-reflogs", "--count", "refs/stash")
 }
 
 func (g *git) hasStaging() bool {
@@ -240,16 +245,3 @@ func groupDict(pattern *regexp.Regexp, haystack string) map[string]string {
 	}
 	return result
 }
-
-func numberOfLinesInString(s string) int {
-	n := 0
-	for _, r := range s {
-		if r == '\n' {
-			n++
-		}
-	}
-	if len(s) > 0 && !strings.HasSuffix(s, "\n") {
-		n++
-	}
-	return n
-}
diff --git a/segment_git_test.go b/segment_git_test.go
index 5ef792d3..6319d85b 100755
--- a/segment_git_test.go
+++ b/segment_git_test.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"math/rand"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -186,9 +185,9 @@ func TestGetGitHEADContextCherryPickOnTag(t *testing.T) {
 }
 
 func TestGetStashContextZeroEntries(t *testing.T) {
-	want := 0
+	want := ""
 	env := new(MockedEnvironment)
-	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "stash", "list"}).Return("")
+	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "rev-list", "--walk-reflogs", "--count", "refs/stash"}).Return("")
 	g := &git{
 		env: env,
 	}
@@ -197,24 +196,9 @@ func TestGetStashContextZeroEntries(t *testing.T) {
 }
 
 func TestGetStashContextMultipleEntries(t *testing.T) {
-	want := rand.Intn(100)
-	var response string
-	for i := 0; i < want; i++ {
-		response += "I'm a stash entry\n"
-	}
+	want := "2"
 	env := new(MockedEnvironment)
-	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "stash", "list"}).Return(response)
-	g := &git{
-		env: env,
-	}
-	got := g.getStashContext()
-	assert.Equal(t, want, got)
-}
-
-func TestGetStashContextOneEntry(t *testing.T) {
-	want := 1
-	env := new(MockedEnvironment)
-	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "stash", "list"}).Return("stash entry")
+	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "rev-list", "--walk-reflogs", "--count", "refs/stash"}).Return("2")
 	g := &git{
 		env: env,
 	}