From 7f3e50236e832a310838f984846959e815f9fb3a Mon Sep 17 00:00:00 2001
From: sairoutine <sairoutine@gmail.com>
Date: Fri, 8 Mar 2019 21:24:07 +0900
Subject: [PATCH] IsTerm flag should not be affected by DisableConsoleColor
 method.

---
 logger.go      | 43 +++++++++++++++++++++++++++----------------
 logger_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 70 insertions(+), 20 deletions(-)

diff --git a/logger.go b/logger.go
index 2ecaed7d..f3e407dd 100644
--- a/logger.go
+++ b/logger.go
@@ -14,17 +14,24 @@ import (
 	"github.com/mattn/go-isatty"
 )
 
+type consoleColorModeValue int
+
+const (
+	autoColor consoleColorModeValue = iota
+	disableColor
+	forceColor
+)
+
 var (
-	green        = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
-	white        = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
-	yellow       = string([]byte{27, 91, 57, 48, 59, 52, 51, 109})
-	red          = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
-	blue         = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
-	magenta      = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
-	cyan         = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
-	reset        = string([]byte{27, 91, 48, 109})
-	disableColor = false
-	forceColor   = false
+	green            = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
+	white            = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
+	yellow           = string([]byte{27, 91, 57, 48, 59, 52, 51, 109})
+	red              = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
+	blue             = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
+	magenta          = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
+	cyan             = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
+	reset            = string([]byte{27, 91, 48, 109})
+	consoleColorMode = autoColor
 )
 
 // LoggerConfig defines the config for Logger middleware.
@@ -115,10 +122,15 @@ func (p *LogFormatterParams) ResetColor() string {
 	return reset
 }
 
+// IsOutputColor indicates whether can colors be outputted to the log.
+func (p *LogFormatterParams) IsOutputColor() bool {
+	return consoleColorMode == forceColor || (consoleColorMode == autoColor && p.IsTerm)
+}
+
 // defaultLogFormatter is the default log format function Logger middleware uses.
 var defaultLogFormatter = func(param LogFormatterParams) string {
 	var statusColor, methodColor, resetColor string
-	if param.IsTerm {
+	if param.IsOutputColor() {
 		statusColor = param.StatusCodeColor()
 		methodColor = param.MethodColor()
 		resetColor = param.ResetColor()
@@ -137,12 +149,12 @@ var defaultLogFormatter = func(param LogFormatterParams) string {
 
 // DisableConsoleColor disables color output in the console.
 func DisableConsoleColor() {
-	disableColor = true
+	consoleColorMode = disableColor
 }
 
 // ForceConsoleColor force color output in the console.
 func ForceConsoleColor() {
-	forceColor = true
+	consoleColorMode = forceColor
 }
 
 // ErrorLogger returns a handlerfunc for any error type.
@@ -199,9 +211,8 @@ func LoggerWithConfig(conf LoggerConfig) HandlerFunc {
 
 	isTerm := true
 
-	if w, ok := out.(*os.File); (!ok ||
-		(os.Getenv("TERM") == "dumb" || (!isatty.IsTerminal(w.Fd()) && !isatty.IsCygwinTerminal(w.Fd()))) ||
-		disableColor) && !forceColor {
+	if w, ok := out.(*os.File); !ok || os.Getenv("TERM") == "dumb" ||
+		(!isatty.IsTerminal(w.Fd()) && !isatty.IsCygwinTerminal(w.Fd())) {
 		isTerm = false
 	}
 
diff --git a/logger_test.go b/logger_test.go
index a2041773..4e160a03 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -297,6 +297,39 @@ func TestResetColor(t *testing.T) {
 	assert.Equal(t, string([]byte{27, 91, 48, 109}), p.ResetColor())
 }
 
+func TestIsOutputColor(t *testing.T) {
+	// test with IsTerm flag true.
+	p := LogFormatterParams{
+		IsTerm: true,
+	}
+
+	consoleColorMode = autoColor
+	assert.Equal(t, true, p.IsOutputColor())
+
+	ForceConsoleColor()
+	assert.Equal(t, true, p.IsOutputColor())
+
+	DisableConsoleColor()
+	assert.Equal(t, false, p.IsOutputColor())
+
+	// test with IsTerm flag false.
+	p = LogFormatterParams{
+		IsTerm: false,
+	}
+
+	consoleColorMode = autoColor
+	assert.Equal(t, false, p.IsOutputColor())
+
+	ForceConsoleColor()
+	assert.Equal(t, true, p.IsOutputColor())
+
+	DisableConsoleColor()
+	assert.Equal(t, false, p.IsOutputColor())
+
+	// reset console color mode.
+	consoleColorMode = autoColor
+}
+
 func TestErrorLogger(t *testing.T) {
 	router := New()
 	router.Use(ErrorLogger())
@@ -359,14 +392,20 @@ func TestLoggerWithConfigSkippingPaths(t *testing.T) {
 
 func TestDisableConsoleColor(t *testing.T) {
 	New()
-	assert.False(t, disableColor)
+	assert.Equal(t, autoColor, consoleColorMode)
 	DisableConsoleColor()
-	assert.True(t, disableColor)
+	assert.Equal(t, disableColor, consoleColorMode)
+
+	// reset console color mode.
+	consoleColorMode = autoColor
 }
 
 func TestForceConsoleColor(t *testing.T) {
 	New()
-	assert.False(t, forceColor)
+	assert.Equal(t, autoColor, consoleColorMode)
 	ForceConsoleColor()
-	assert.True(t, forceColor)
+	assert.Equal(t, forceColor, consoleColorMode)
+
+	// reset console color mode.
+	consoleColorMode = autoColor
 }