From 9a108380de1fa573eb48069eabc5649637843d06 Mon Sep 17 00:00:00 2001
From: MichaelC <michael@thisisbud.com>
Date: Tue, 27 Aug 2019 23:49:38 +0100
Subject: [PATCH] Add the ability to override log level values

---
 json_formatter.go      | 28 +++++++++++++++++++++++++++-
 json_formatter_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/json_formatter.go b/json_formatter.go
index 098a21a..af3455a 100644
--- a/json_formatter.go
+++ b/json_formatter.go
@@ -20,6 +20,17 @@ func (f FieldMap) resolve(key fieldKey) string {
 	return string(key)
 }
 
+// LogLevel map allows customization of the values used to describe log levels
+type LogLevelMap map[Level]string
+
+func (l LogLevelMap) resolve(key Level) string {
+	if k, ok := l[key]; ok {
+		return k
+	}
+
+	return key.String()
+}
+
 // JSONFormatter formats logs into parsable json
 type JSONFormatter struct {
 	// TimestampFormat sets the format used for marshaling timestamps.
@@ -43,6 +54,21 @@ type JSONFormatter struct {
 	// }
 	FieldMap FieldMap
 
+	// LogLevelMap allows users to customize the values used for the various log levels
+	// As an example:
+	// formatter := &JSONFormatter{
+	//   	LogLevelMap: LogLevelMap{
+	// 		 PanicLevel:  1,
+	// 		 FatalLevel:  2,
+	// 		 ErrorLevel:  3,
+	// 		 WarnLevel:   4,
+	// 		 InfoLevel:   5,
+	// 		 DebugLevel:  6,
+	// 		 TraceLevel:  7,
+	//    },
+	// }
+	LogLevelMap LogLevelMap
+
 	// CallerPrettyfier can be set by the user to modify the content
 	// of the function and file keys in the json data when ReportCaller is
 	// activated. If any of the returned value is the empty string the
@@ -87,7 +113,7 @@ func (f *JSONFormatter) Format(entry *Entry) ([]byte, error) {
 		data[f.FieldMap.resolve(FieldKeyTime)] = entry.Time.Format(timestampFormat)
 	}
 	data[f.FieldMap.resolve(FieldKeyMsg)] = entry.Message
-	data[f.FieldMap.resolve(FieldKeyLevel)] = entry.Level.String()
+	data[f.FieldMap.resolve(FieldKeyLevel)] = f.LogLevelMap.resolve(entry.Level)
 	if entry.HasCaller() {
 		funcVal := entry.Caller.Function
 		fileVal := fmt.Sprintf("%s:%d", entry.Caller.File, entry.Caller.Line)
diff --git a/json_formatter_test.go b/json_formatter_test.go
index 695c36e..5ea6451 100644
--- a/json_formatter_test.go
+++ b/json_formatter_test.go
@@ -234,6 +234,33 @@ func TestJSONMessageKey(t *testing.T) {
 	}
 }
 
+func TestLogLevelMap(t *testing.T) {
+	formatter := &JSONFormatter{
+		LogLevelMap: LogLevelMap{
+			WarnLevel: "WARNING",
+		},
+	}
+
+	b, err := formatter.Format(&Entry{Message: "oh hai", Level: WarnLevel})
+	if err != nil {
+		t.Fatal("Unable to format entry: ", err)
+	}
+	s := string(b)
+	if !(strings.Contains(s, "WARNING")) {
+		t.Fatal("Expected to override log level field value")
+	}
+
+	entry := make(map[string]interface{})
+	err = json.Unmarshal(b, &entry)
+	if err != nil {
+		t.Fatal("Unable to unmarshal formatted entry: ", err)
+	}
+
+	if entry["level"] != "WARNING" {
+		t.Fatal("level field not set correctly")
+	}
+}
+
 func TestJSONLevelKey(t *testing.T) {
 	formatter := &JSONFormatter{
 		FieldMap: FieldMap{