From 8f1c64ba3092ddf4f5190323483a9a8344d1a745 Mon Sep 17 00:00:00 2001
From: Mauricio Trajano <mauriciot1993@gmail.com>
Date: Thu, 26 Dec 2024 20:30:44 -0500
Subject: [PATCH] Add ability to configure branch color patterns

---
 docs/Config.md                   |  7 ++++++
 pkg/config/user_config.go        |  4 +++-
 pkg/gui/gui.go                   |  8 ++++++-
 pkg/gui/presentation/branches.go | 41 ++++++++++++++++++++++++++------
 4 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/docs/Config.md b/docs/Config.md
index c64e56d2edf..179790a9f27 100644
--- a/docs/Config.md
+++ b/docs/Config.md
@@ -812,8 +812,15 @@ You can customize the color of branches based on the branch prefix:
 
 ```yaml
 gui:
+  # NOTE: that this configuration will be deprecated in favor of using branchColorPatterns below
   branchColors:
     'docs': '#11aaff' # use a light blue for branches beginning with 'docs/'
+
+  # alternatively you can use a regex pattern as your coloring rules
+  # NOTE: this configuration overwrites the one above, if you would like to set a similar rule see the example below
+  branchColorPatterns:
+    'docs/.+': '#11aaff' # similar to the previous configuration above, setting branches that begin with 'docs/'
+    'ISSUE-\d+': '#ff5733' # use a bright orange for branches beginning with 'ISSUE-'
 ```
 
 ## Example Coloring
diff --git a/pkg/config/user_config.go b/pkg/config/user_config.go
index 36dc02a68a7..36230ba85d7 100644
--- a/pkg/config/user_config.go
+++ b/pkg/config/user_config.go
@@ -52,7 +52,9 @@ type GuiConfig struct {
 	// See https://github.com/jesseduffield/lazygit/blob/master/docs/Config.md#custom-author-color
 	AuthorColors map[string]string `yaml:"authorColors"`
 	// See https://github.com/jesseduffield/lazygit/blob/master/docs/Config.md#custom-branch-color
-	BranchColors map[string]string `yaml:"branchColors"`
+	// NOTE: BranchColors is being deprecated in favor of BranchColorPatterns
+	BranchColors        map[string]string `yaml:"branchColors"`
+	BranchColorPatterns map[string]string `yaml:"branchColorPatterns"`
 	// The number of lines you scroll by when scrolling the main window
 	ScrollHeight int `yaml:"scrollHeight" jsonschema:"minimum=1"`
 	// If true, allow scrolling past the bottom of the content in the main window
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 1faeb85e4c0..2276b7c473e 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -454,7 +454,13 @@ func (gui *Gui) onUserConfigLoaded() error {
 	} else if userConfig.Gui.ShowIcons {
 		icons.SetNerdFontsVersion("2")
 	}
-	presentation.SetCustomBranches(userConfig.Gui.BranchColors)
+
+	if userConfig.Gui.BranchColorPatterns != nil {
+		presentation.SetCustomBranches(userConfig.Gui.BranchColorPatterns, true)
+	} else {
+		// The alternative is to match on branch types with the branchColors config which will be deprecated in the future
+		presentation.SetCustomBranches(userConfig.Gui.BranchColors, false)
+	}
 
 	return nil
 }
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index b75dfc95b72..d91747145bb 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -2,6 +2,7 @@ package presentation
 
 import (
 	"fmt"
+	"regexp"
 	"strings"
 	"time"
 
@@ -18,7 +19,12 @@ import (
 	"github.com/samber/lo"
 )
 
-var branchPrefixColorCache = make(map[string]style.TextStyle)
+type colorMatcher struct {
+	patterns map[string]style.TextStyle
+	isRegex  bool // NOTE: this value is needed only until branchColors is deprecated and only regex color patterns are used
+}
+
+var branchPrefixColorCache *colorMatcher
 
 func GetBranchListDisplayStrings(
 	branches []*models.Branch,
@@ -125,12 +131,12 @@ func getBranchDisplayStrings(
 
 // GetBranchTextStyle branch color
 func GetBranchTextStyle(name string) style.TextStyle {
-	branchType := strings.Split(name, "/")[0]
-
-	if value, ok := branchPrefixColorCache[branchType]; ok {
-		return value
+	if style, ok := branchPrefixColorCache.match(name); ok {
+		return *style
 	}
 
+	// Default colors for common branch types
+	branchType := strings.Split(name, "/")[0]
 	switch branchType {
 	case "feature":
 		return style.FgGreen
@@ -143,6 +149,24 @@ func GetBranchTextStyle(name string) style.TextStyle {
 	}
 }
 
+func (m *colorMatcher) match(name string) (*style.TextStyle, bool) {
+	if m.isRegex {
+		for pattern, style := range m.patterns {
+			if matched, _ := regexp.MatchString("^"+pattern+"$", name); matched {
+				return &style, true
+			}
+		}
+	} else {
+		// old behavior using the deprecated branchColors behavior matching on branch type
+		branchType := strings.Split(name, "/")[0]
+		if value, ok := m.patterns[branchType]; ok {
+			return &value, true
+		}
+	}
+
+	return nil, false
+}
+
 func BranchStatus(
 	branch *models.Branch,
 	itemOperation types.ItemOperation,
@@ -189,6 +213,9 @@ func BranchStatus(
 	return result
 }
 
-func SetCustomBranches(customBranchColors map[string]string) {
-	branchPrefixColorCache = utils.SetCustomColors(customBranchColors)
+func SetCustomBranches(customBranchColors map[string]string, isRegex bool) {
+	branchPrefixColorCache = &colorMatcher{
+		patterns: utils.SetCustomColors(customBranchColors),
+		isRegex:  isRegex,
+	}
 }