From c8f3704d51c95bcfbf45f63b137cd5940a45a9ac Mon Sep 17 00:00:00 2001
From: Randy <randywijnants88@gmail.com>
Date: Mon, 23 Oct 2023 14:36:24 +0200
Subject: [PATCH] Add custom level color functionality to TextFormatter

---
 README.md         | 34 +++++++++++++++++++++++++++
 text_formatter.go | 60 +++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 87 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index d1d4a85..e786bd1 100644
--- a/README.md
+++ b/README.md
@@ -319,6 +319,40 @@ environment if your application has that.
 
 Note: If you want different log levels for global (`log.SetLevel(...)`) and syslog logging, please check the [syslog hook README](hooks/syslog/README.md#different-log-levels-for-local-and-remote-logging).
 
+#### Level colors
+
+The `TextFormatter` level colors itself can be customized as follows:
+
+Individual overrides:
+
+```go
+	log.SetFormatter(&log.TextFormatter{
+		LevelColors: log.DefaultLevelColors(&log.LevelColors{
+			Debug:   31, // Red
+      Default: 36, // Blue
+		}),
+		ForceColors: true, // Optional
+	})
+```
+
+Or you can provide a complete `log.LevelColors` struct: 
+
+```go
+	log.SetFormatter(&log.TextFormatter{
+		LevelColors: &log.LevelColors{
+			Trace:   37, // Gray
+			Debug:   37, // Gray
+			Warn:    33, // Yellow
+			Error:   31, // Red
+			Fatal:   31, // Red
+			Panic:   31, // Red
+			Info:    36, // Blue
+			Default: 36, // Blue
+		},
+		ForceColors: true, // Optional
+	})
+```
+
 #### Entries
 
 Besides the fields added with `WithField` or `WithFields` some fields are
diff --git a/text_formatter.go b/text_formatter.go
index be2c6ef..82d3636 100644
--- a/text_formatter.go
+++ b/text_formatter.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"os"
+	"reflect"
 	"runtime"
 	"sort"
 	"strconv"
@@ -26,8 +27,43 @@ func init() {
 	baseTimestamp = time.Now()
 }
 
+type LevelColors struct {
+	Trace, Debug, Warn, Error, Fatal, Panic, Info, Default int
+}
+
+func DefaultLevelColors(custom *LevelColors) *LevelColors {
+	levelColors := &LevelColors{
+		Trace:   gray,
+		Debug:   gray,
+		Warn:    yellow,
+		Error:   red,
+		Fatal:   red,
+		Panic:   red,
+		Info:    blue,
+		Default: blue,
+	}
+
+	if custom != nil {
+		dstVal := reflect.ValueOf(levelColors).Elem()
+		srcVal := reflect.ValueOf(custom).Elem()
+		for i := 0; i < dstVal.NumField(); i++ {
+			dstField := dstVal.Field(i)
+			srcField := srcVal.Field(i)
+
+			if dstField.Kind() == reflect.Int && srcField.Int() != 0 {
+				dstField.SetInt(srcField.Int())
+			}
+		}
+	}
+
+	return levelColors
+}
+
 // TextFormatter formats logs into text
 type TextFormatter struct {
+	// Set level colors
+	LevelColors *LevelColors
+
 	// Set to true to bypass checking for a TTY before outputting colors.
 	ForceColors bool
 
@@ -230,18 +266,28 @@ func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {
 }
 
 func (f *TextFormatter) printColored(b *bytes.Buffer, entry *Entry, keys []string, data Fields, timestampFormat string) {
+	if f.LevelColors == nil {
+		f.LevelColors = DefaultLevelColors(nil)
+	}
+
 	var levelColor int
 	switch entry.Level {
-	case DebugLevel, TraceLevel:
-		levelColor = gray
+	case DebugLevel:
+		levelColor = f.LevelColors.Debug
+	case TraceLevel:
+		levelColor = f.LevelColors.Trace
 	case WarnLevel:
-		levelColor = yellow
-	case ErrorLevel, FatalLevel, PanicLevel:
-		levelColor = red
+		levelColor = f.LevelColors.Warn
+	case ErrorLevel:
+		levelColor = f.LevelColors.Error
+	case FatalLevel:
+		levelColor = f.LevelColors.Fatal
+	case PanicLevel:
+		levelColor = f.LevelColors.Panic
 	case InfoLevel:
-		levelColor = blue
+		levelColor = f.LevelColors.Info
 	default:
-		levelColor = blue
+		levelColor = f.LevelColors.Default
 	}
 
 	levelText := strings.ToUpper(entry.Level.String())