common/loggers: Add color to ERROR and WARN

Fixes #4414

2 files changed, 65 insertions, 2 deletions

diff --git a/common/loggers/loggers.go b/common/loggers/loggers.go
index 428a9f79c..37f50a6c0 100644
--- a/common/loggers/loggers.go
+++ b/common/loggers/loggers.go
@@ -80,7 +80,10 @@ func NewErrorLogger() *Logger {
 	return newBasicLogger(jww.LevelError)
 }
 
-var ansiColorRe = regexp.MustCompile("(?s)\\033\\[\\d*(;\\d*)*m")
+var (
+	ansiColorRe = regexp.MustCompile("(?s)\\033\\[\\d*(;\\d*)*m")
+	errorRe     = regexp.MustCompile("(ERROR|FATAL|WARN)")
+)
 
 type ansiCleaner struct {
 	w io.Writer
@@ -90,12 +93,40 @@ func (a ansiCleaner) Write(p []byte) (n int, err error) {
 	return a.w.Write(ansiColorRe.ReplaceAll(p, []byte("")))
 }
 
+type labelColorizer struct {
+	w io.Writer
+}
+
+func (a labelColorizer) Write(p []byte) (n int, err error) {
+	replaced := errorRe.ReplaceAllStringFunc(string(p), func(m string) string {
+		switch m {
+		case "ERROR", "FATAL":
+			return terminal.Error(m)
+		case "WARN":
+			return terminal.Warning(m)
+		default:
+			return m
+		}
+	})
+	// io.MultiWriter will abort if we return a bigger write count than input
+	// bytes, so we lie a little.
+	_, err = a.w.Write([]byte(replaced))
+	return len(p), err
+
+}
+
 func newLogger(stdoutThreshold, logThreshold jww.Threshold, outHandle, logHandle io.Writer, saveErrors bool) *Logger {
+	isTerm := terminal.IsTerminal(os.Stdout)
 	errorCounter := &jww.Counter{}
-	if logHandle != ioutil.Discard && terminal.IsTerminal(os.Stdout) {
+	if logHandle != ioutil.Discard && isTerm {
 		// Remove any Ansi coloring from log output
 		logHandle = ansiCleaner{w: logHandle}
 	}
+
+	if isTerm {
+		outHandle = labelColorizer{w: outHandle}
+	}
+
 	listeners := []jww.LogListener{jww.LogCounter(errorCounter, jww.LevelError)}
 	var errorBuff *bytes.Buffer
 	if saveErrors {
diff --git a/common/loggers/loggers_test.go b/common/loggers/loggers_test.go
new file mode 100644
index 000000000..3737ddc68
--- /dev/null
+++ b/common/loggers/loggers_test.go
@@ -0,0 +1,32 @@
+// Copyright 2018 The Hugo Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package loggers
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLogger(t *testing.T) {
+	assert := require.New(t)
+	l := NewWarningLogger()
+
+	l.ERROR.Println("One error")
+	l.ERROR.Println("Two error")
+	l.WARN.Println("A warning")
+
+	assert.Equal(uint64(2), l.ErrorCounter.Count())
+
+}