From a05b03a5aada46acdcf24bcff0e4670a3795449d Mon Sep 17 00:00:00 2001
From: Laurent Demailly <ldemailly@gmail.com>
Date: Mon, 31 Jul 2023 18:27:06 -0700
Subject: [PATCH] Allow caller code to use Colors directly without if  (#40)

* Allow caller code to use Colors directly without if and they'll get reset to empty string if color mode is disabled (and vice versa)

* linter fix + added test

* better doc
---
 README.md          |  2 ++
 console_logging.go | 16 +++++++++++++++-
 logger_test.go     | 23 +++++++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 93e7218..9f477a5 100644
--- a/README.md
+++ b/README.md
@@ -72,3 +72,5 @@ When on console:
 ![Example console color output](color.png)
 
 JSON formatted logs can also be converted back to text later/after capture and similarly colorized using [fortio.org/logc](https://github.com/fortio/logc#logc)
+
+The `log.Colors` can be used by callers and they'll be empty string when not in color mode, and the ansi escape codes otherwise.
diff --git a/console_logging.go b/console_logging.go
index 010114e..f20dffd 100644
--- a/console_logging.go
+++ b/console_logging.go
@@ -41,7 +41,9 @@ var (
 	// these should really be constants but go doesn't have constant structs, arrays etc...
 
 	// ANSI color codes.
-	Colors = color{
+	// This isn't meant to be used directly and is here only to document the names of the struct.
+	// Use the Colors variable instead.
+	ANSIColors = color{
 		Reset:     "\033[0m",
 		Red:       "\033[31m",
 		Green:     "\033[32m",
@@ -55,6 +57,11 @@ var (
 		DarkGray:  "\033[90m",
 	}
 
+	// ANSI color codes or empty depending on ColorMode.
+	// These will be reset to empty string if color is disabled (see ColorMode() and SetColorMode()).
+	// Start with actual colors, will be reset to empty if color is disabled.
+	Colors = ANSIColors
+
 	// Mapping of log levels to color.
 	LevelToColor = []string{
 		Colors.Gray,
@@ -81,8 +88,15 @@ func ConsoleLogging() bool {
 
 // SetColorMode computes whether we currently should be using color text mode or not.
 // Need to be reset if config changes (but is already automatically re evaluated when calling SetOutput()).
+// It will reset the Colors variable to either be the actual escape sequences or empty strings (when
+// color is disabled).
 func SetColorMode() {
 	Color = ColorMode()
+	if Color {
+		Colors = ANSIColors
+	} else {
+		Colors = color{}
+	}
 }
 
 // ColorMode returns true if we should be using color text mode, which is either because it's
diff --git a/logger_test.go b/logger_test.go
index 489d3cd..750dcee 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -158,8 +158,31 @@ func TestColorMode(t *testing.T) {
 	if cs == "" {
 		t.Errorf("got empty color timestamp")
 	}
+	if Colors.Green == "" {
+		t.Errorf("expected to have green color not empty when in color mode")
+	}
+	prevGreen := Colors.Green
+	// turn off color mode
+	Config.ForceColor = false
+	SetColorMode()
+	if Color {
+		t.Errorf("expected to not be in color mode after SetColorMode() and forcecolor false")
+	}
+	if Colors.Green != "" {
+		t.Errorf("expected to have green color empty when not color mode, got %q", Colors.Green)
+	}
+	// Show one can mutate/change/tweak colors
+	ANSIColors.Green = "foo"
+	Config.ForceColor = true
+	SetColorMode()
+	if Colors.Green != "foo" {
+		t.Errorf("expected to have green color preserved, got %q", Colors.Green)
+	}
+	// put it back to real green for other tests
+	ANSIColors.Green = prevGreen
 	// Reset for other/further tests
 	Config.ForceColor = false
+	SetColorMode()
 }
 
 func TestSetLevel(t *testing.T) {