From 93bba0ad75e2996191e16c1f33bb88d78c65f284 Mon Sep 17 00:00:00 2001
From: Mauricio Trajano <mauriciot1993@gmail.com>
Date: Wed, 1 Jan 2025 18:49:25 -0500
Subject: [PATCH] Overwrite branchColors config

---
 docs/Config.md                   |  7 +++-
 pkg/config/user_config.go        |  3 +-
 pkg/gui/gui.go                   |  7 +++-
 pkg/gui/presentation/branches.go | 57 ++++++++++++++++++--------------
 4 files changed, 47 insertions(+), 27 deletions(-)

diff --git a/docs/Config.md b/docs/Config.md
index 7c9d142c832..dbefb760d9f 100644
--- a/docs/Config.md
+++ b/docs/Config.md
@@ -814,7 +814,12 @@ You can customize the color of branches based on the branch prefix:
 gui:
   branchColors:
     'docs': '#11aaff' # use a light blue for branches beginning with 'docs/'
-    'ISSUE-*': '#ff5733' # use a bright orange for branches beginning with 'ISSUE-'
+
+  # 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..69fc01f94c4 100644
--- a/pkg/config/user_config.go
+++ b/pkg/config/user_config.go
@@ -52,7 +52,8 @@ 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"`
+	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..cc62323e41f 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -454,7 +454,12 @@ 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 {
+		presentation.SetCustomBranches(userConfig.Gui.BranchColors, false)
+	}
 
 	return nil
 }
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index bc4e7e6bb18..096e745f8bc 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
+}
+
+var branchPrefixColorCache *colorMatcher
 
 func GetBranchListDisplayStrings(
 	branches []*models.Branch,
@@ -123,31 +129,14 @@ func getBranchDisplayStrings(
 	return res
 }
 
-func matchColorPrefix(name string) (style.TextStyle, bool) {
-	for possiblePattern, color := range branchPrefixColorCache {
-		if strings.Contains(possiblePattern, "*") {
-			prefix := strings.TrimSuffix(possiblePattern, "*")
-			if strings.HasPrefix(name, prefix) {
-				return color, true
-			}
-		}
-	}
-
-	return style.TextStyle{}, false
-}
-
 // GetBranchTextStyle branch color
 func GetBranchTextStyle(name string) style.TextStyle {
-	branchType := strings.Split(name, "/")[0]
-
-	if value, ok := branchPrefixColorCache[branchType]; ok {
-		return value
-	}
-
-	if value, ok := matchColorPrefix(name); 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
@@ -160,6 +149,23 @@ 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 {
+		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,
@@ -206,6 +212,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,
+	}
 }