From 9227990a3b2a62e3cedba51225b03f63cc9f3af5 Mon Sep 17 00:00:00 2001
From: SuperQ <superq@gmail.com>
Date: Sat, 25 Feb 2023 12:57:58 +0100
Subject: [PATCH] Add config / flag for color.

Signed-off-by: SuperQ <superq@gmail.com>
---
 promlog/flag/flag.go | 10 ++++++++
 promlog/log.go       | 59 ++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/promlog/flag/flag.go b/promlog/flag/flag.go
index ec55008b..9a0c831e 100644
--- a/promlog/flag/flag.go
+++ b/promlog/flag/flag.go
@@ -18,6 +18,12 @@ import (
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 )
 
+// ColorFlagName is the canonical flag name to enable color logging.
+const ColorFlagName = "log.color"
+
+// ColorFlagHelp is the help description for the log.color flag.
+const ColorFlagHelp = "Enable colors in the console log. Use --no-log.color to disable."
+
 // LevelFlagName is the canonical flag name to configure the allowed log level
 // within Prometheus projects.
 const LevelFlagName = "log.level"
@@ -35,6 +41,10 @@ const FormatFlagHelp = "Output format of log messages. One of: [logfmt, json]"
 // AddFlags adds the flags used by this package to the Kingpin application.
 // To use the default Kingpin application, call AddFlags(kingpin.CommandLine)
 func AddFlags(a *kingpin.Application, config *promlog.Config) {
+	config.Color = &promlog.Color{}
+	a.Flag(ColorFlagName, ColorFlagHelp).
+		Default("true").SetValue(config.Color)
+
 	config.Level = &promlog.AllowedLevel{}
 	a.Flag(LevelFlagName, LevelFlagHelp).
 		Default("info").SetValue(config.Level)
diff --git a/promlog/log.go b/promlog/log.go
index 0751b396..63908e3b 100644
--- a/promlog/log.go
+++ b/promlog/log.go
@@ -19,6 +19,7 @@ package promlog
 import (
 	"fmt"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
@@ -37,6 +38,33 @@ var (
 	)
 )
 
+// Color is a settable boolean for controlling color output.
+type Color struct {
+	s       string
+	enabled bool
+}
+
+func (c *Color) Set(s string) error {
+	switch s {
+	case "true":
+		c.enabled = true
+	case "false":
+		c.enabled = false
+	default:
+		return fmt.Errorf("unrecognized boolean %q", s)
+	}
+	c.s = s
+	return nil
+}
+
+func (c *Color) Enabled() bool {
+	return c.enabled
+}
+
+func (c *Color) String() string {
+	return strconv.FormatBool(c.enabled)
+}
+
 // AllowedLevel is a settable identifier for the minimum level a log entry
 // must be have.
 type AllowedLevel struct {
@@ -124,6 +152,7 @@ func (f *AllowedFormat) Set(s string) error {
 
 // Config is a struct containing configurable settings for the logger
 type Config struct {
+	Color  *Color
 	Level  *AllowedLevel
 	Format *AllowedFormat
 }
@@ -131,14 +160,22 @@ type Config struct {
 // New returns a new leveled oklog logger. Each logged line will be annotated
 // with a timestamp. The output always goes to stderr.
 func New(config *Config) log.Logger {
+	if config.Color == nil {
+		config.Color = &Color{}
+		config.Color.Set("true")
+	}
 	var l log.Logger
 	syncWriter := log.NewSyncWriter(os.Stderr)
 	if config.Format != nil && config.Format.s == "json" {
 		l = log.NewJSONLogger(syncWriter)
 	} else {
-		// Returns a new logger with color logging capabilites if we're in a terminal, otherwise we
-		// just get a standard go-kit logger.
-		l = term.NewLogger(syncWriter, log.NewLogfmtLogger, colorFn)
+		if config.Color.Enabled() {
+			// Returns a new logger with color logging capabilites if we're in a terminal, otherwise we
+			// just get a standard go-kit logger.
+			l = term.NewLogger(syncWriter, log.NewLogfmtLogger, colorFn)
+		} else {
+			l = log.NewJSONLogger(syncWriter)
+		}
 	}
 
 	if config.Level != nil {
@@ -154,11 +191,23 @@ func New(config *Config) log.Logger {
 // with a timestamp. The output always goes to stderr. Some properties can be
 // changed, like the level.
 func NewDynamic(config *Config) *logger {
+	if config.Color == nil {
+		config.Color = &Color{}
+		config.Color.Set("true")
+	}
 	var l log.Logger
+	syncWriter := log.NewSyncWriter(os.Stderr)
+
 	if config.Format != nil && config.Format.s == "json" {
-		l = log.NewJSONLogger(log.NewSyncWriter(os.Stderr))
+		l = log.NewJSONLogger(syncWriter)
 	} else {
-		l = log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
+		if config.Color.Enabled() {
+			// Returns a new logger with color logging capabilites if we're in a terminal, otherwise we
+			// just get a standard go-kit logger.
+			l = term.NewLogger(syncWriter, log.NewLogfmtLogger, colorFn)
+		} else {
+			l = log.NewJSONLogger(syncWriter)
+		}
 	}
 
 	lo := &logger{