From 4210d816c75916f345e260e8ce57969b21bb2f16 Mon Sep 17 00:00:00 2001
From: Pieter Claerhout <pieter@yellowduck.be>
Date: Sat, 26 Oct 2019 09:55:23 +0200
Subject: [PATCH] Add support for printing log messages in color

Fixes #6
---
 README.md          |  1 +
 go.mod             |  3 +++
 go.sum             |  2 ++
 logger.go          |  3 +++
 logger_internal.go | 16 +++++++++++++++-
 logger_test.go     |  1 +
 6 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 98c8bfb..11696b6 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,7 @@ func main() {
     log.DebugMode = true
     log.DebugSQLMode = true
     log.PrintTimestamp = true
+    log.PrintColors = true
     log.TimeFormat = "2006-01-02 15:04:05.000"
 
     myVar := map[string]string{"hello": "world"}
diff --git a/go.mod b/go.mod
index 6520750..500ae84 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,10 @@ module github.com/pieterclaerhout/go-log
 go 1.13
 
 require (
+	github.com/fatih/color v1.7.0
 	github.com/go-errors/errors v1.0.1
+	github.com/mattn/go-colorable v0.1.4 // indirect
+	github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b
 	github.com/pieterclaerhout/go-formatter v1.0.2
 	github.com/pkg/errors v0.8.1
 	github.com/sanity-io/litter v1.1.0
diff --git a/go.sum b/go.sum
index cd275d8..181a380 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,7 @@
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
+github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/go-errors/errors v1.0.1 h1:LUHzmkK3GUKUrL/1gfBUxAHzcev3apQlezX/+O7ma6w=
 github.com/go-errors/errors v1.0.1/go.mod h1:f4zRHt4oKfwPJE5k8C9vpYG+aDHdBFUsgrm6/TyX73Q=
 github.com/mattn/go-colorable v0.1.4 h1:snbPLB8fVfU9iwbbo30TPtbLRzwWu6aJS6Xh4eaaviA=
diff --git a/logger.go b/logger.go
index 70b5da1..3b4d42d 100644
--- a/logger.go
+++ b/logger.go
@@ -16,6 +16,9 @@ import (
 // PrintTimestamp indicates if the log messages should include a timestamp or not
 var PrintTimestamp = false
 
+// PrintColors indicates if the messages should be printed in color or not
+var PrintColors = false
+
 // DebugMode indicates if debug information should be printed or not
 //
 // If the environment variable called DEBUG is set to 1, this will default to true.
diff --git a/logger_internal.go b/logger_internal.go
index cadcf55..6116dde 100644
--- a/logger_internal.go
+++ b/logger_internal.go
@@ -6,10 +6,20 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/fatih/color"
 )
 
 var logMutex = &sync.Mutex{}
 
+var colors = map[string]*color.Color{
+	"DEBUG": color.New(color.FgHiBlack),
+	"INFO ": color.New(color.FgHiGreen),
+	"WARN ": color.New(color.FgHiYellow),
+	"ERROR": color.New(color.FgHiRed),
+	"FATAL": color.New(color.FgHiRed),
+}
+
 func init() {
 	TimeZone, _ = time.LoadLocation("Europe/Brussels")
 	DebugMode = os.Getenv("DEBUG") == "1"
@@ -50,6 +60,11 @@ func printMessage(level string, message string) {
 		formattedTime := tstamp.Format(TimeFormat)
 		message = formattedTime + " | " + level + " | " + message
 	}
+	if PrintColors {
+		if color, ok := colors[level]; ok {
+			message = color.Sprint(message)
+		}
+	}
 
 	w := Stdout
 	if level == "ERROR" || level == "FATAL" {
@@ -57,7 +72,6 @@ func printMessage(level string, message string) {
 	}
 
 	w.Write([]byte(message + "\n"))
-	// fmt.Fprint(w, message+"\n")
 
 	logMutex.Unlock()
 
diff --git a/logger_test.go b/logger_test.go
index d8e66e6..941453d 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -609,6 +609,7 @@ func TestCheckError(t *testing.T) {
 
 func resetLogConfig() {
 	log.PrintTimestamp = true
+	log.PrintColors = true
 	log.DebugMode = false
 	log.DebugSQLMode = false
 	log.TimeZone, _ = time.LoadLocation("Europe/Brussels")