From 9adf4a190864fd0c6a3b00ef46f23e93b0e6a5f5 Mon Sep 17 00:00:00 2001
From: Matt Cles <matthew@cles.dev>
Date: Mon, 31 Jan 2022 21:23:55 -0800
Subject: [PATCH] Add shared function for loading map of custom colors

---
 pkg/gui/gui.go                          |  4 ++--
 pkg/gui/presentation/authors/authors.go |  5 +----
 pkg/gui/presentation/branches.go        | 13 +++++--------
 pkg/utils/color.go                      | 12 ++++++++++++
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index c0cb9692c..e5bcd3c88 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -24,8 +24,8 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gui/modes/cherrypicking"
 	"github.com/jesseduffield/lazygit/pkg/gui/modes/diffing"
 	"github.com/jesseduffield/lazygit/pkg/gui/modes/filtering"
-	"github.com/jesseduffield/lazygit/pkg/gui/presentation/authors"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
+	"github.com/jesseduffield/lazygit/pkg/gui/presentation/authors"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation/graph"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
@@ -493,7 +493,7 @@ func NewGui(
 	gui.PopupHandler = &RealPopupHandler{gui: gui}
 
 	authors.SetCustomAuthors(gui.UserConfig.Gui.AuthorColors)
-	presentation.SetCustomBranchColors(gui.UserConfig.Gui.BranchColors)
+	presentation.SetCustomBranches(gui.UserConfig.Gui.BranchColors)
 
 	return gui, nil
 }
diff --git a/pkg/gui/presentation/authors/authors.go b/pkg/gui/presentation/authors/authors.go
index 7924ba369..d4eb8e465 100644
--- a/pkg/gui/presentation/authors/authors.go
+++ b/pkg/gui/presentation/authors/authors.go
@@ -114,8 +114,5 @@ func getFirstRune(str string) rune {
 }
 
 func SetCustomAuthors(customAuthorColors map[string]string) {
-	for authorName, colorSequence := range customAuthorColors {
-		style := style.New().SetFg(style.NewRGBColor(color.HEX(colorSequence, false)))
-		authorStyleCache[authorName] = style
-	}
+	authorStyleCache = utils.SetCustomColors(customAuthorColors)
 }
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index 50a4ef094..e31c823ce 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -4,13 +4,13 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/gookit/color"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/theme"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
-var branchPrefixColors = make(map[string]style.TextStyle)
+var branchPrefixColorCache = make(map[string]style.TextStyle)
 
 func GetBranchListDisplayStrings(branches []*models.Branch, fullDescription bool, diffName string) [][]string {
 	lines := make([][]string, len(branches))
@@ -61,7 +61,7 @@ func getBranchDisplayStrings(b *models.Branch, fullDescription bool, diffed bool
 func GetBranchTextStyle(name string) style.TextStyle {
 	branchType := strings.Split(name, "/")[0]
 
-	if value, ok := branchPrefixColors[branchType]; ok {
+	if value, ok := branchPrefixColorCache[branchType]; ok {
 		return value
 	}
 
@@ -92,9 +92,6 @@ func BranchStatus(branch *models.Branch) string {
 	return fmt.Sprintf("↑%s↓%s", branch.Pushables, branch.Pullables)
 }
 
-func SetCustomBranchColors(customBranchColors map[string]string) {
-	for branchPrefix, colorSequence := range customBranchColors {
-		style := style.New().SetFg(style.NewRGBColor(color.HEX(colorSequence, false)))
-		branchPrefixColors[branchPrefix] = style
-	}
+func SetCustomBranches(customBranchColors map[string]string) {
+	branchPrefixColorCache = utils.SetCustomColors(customBranchColors)
 }
diff --git a/pkg/utils/color.go b/pkg/utils/color.go
index 76e4c41c1..37c60179a 100644
--- a/pkg/utils/color.go
+++ b/pkg/utils/color.go
@@ -3,6 +3,9 @@ package utils
 import (
 	"regexp"
 	"sync"
+
+	"github.com/gookit/color"
+	"github.com/jesseduffield/lazygit/pkg/gui/style"
 )
 
 var decoloriseCache = make(map[string]string)
@@ -48,3 +51,12 @@ func IsValidHexValue(v string) bool {
 
 	return true
 }
+
+func SetCustomColors(customColors map[string]string) map[string]style.TextStyle {
+	colors := make(map[string]style.TextStyle)
+	for key, colorSequence := range customColors {
+		style := style.New().SetFg(style.NewRGBColor(color.HEX(colorSequence, false)))
+		colors[key] = style
+	}
+	return colors
+}