From 4f09b8fd5e47dbccfe64ef91a464194b0e8ed053 Mon Sep 17 00:00:00 2001
From: Bryan Conn <30739012+bconn98@users.noreply.github.com>
Date: Mon, 22 Jan 2024 23:04:30 -0500
Subject: [PATCH] Add windows and *nix based env vars to control color output
 in logs (#335)

---
 Cargo.toml                   |   1 +
 src/encode/writer/console.rs | 101 +++++++++++++++++++++++++++++------
 tests/color_control.rs       |  24 +++++++++
 3 files changed, 111 insertions(+), 15 deletions(-)
 create mode 100644 tests/color_control.rs

diff --git a/Cargo.toml b/Cargo.toml
index 629999aa..18b1a81d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -71,6 +71,7 @@ parking_lot = { version = "0.12.0", optional = true }
 thiserror = "1.0.15"
 anyhow = "1.0.28"
 derivative = "2.2"
+once_cell = "1.17.1"
 
 [target.'cfg(windows)'.dependencies]
 winapi = { version = "0.3", optional = true, features = ["handleapi", "minwindef", "processenv", "winbase", "wincon"] }
diff --git a/src/encode/writer/console.rs b/src/encode/writer/console.rs
index fb671850..6209fb97 100644
--- a/src/encode/writer/console.rs
+++ b/src/encode/writer/console.rs
@@ -5,6 +5,42 @@
 use std::{fmt, io};
 
 use crate::encode::{self, Style};
+use once_cell::sync::Lazy;
+
+static COLOR_MODE: Lazy<ColorMode> = Lazy::new(|| {
+    let no_color = std::env::var("NO_COLOR")
+        .map(|var| var != "0")
+        .unwrap_or(false);
+    let clicolor_force = std::env::var("CLICOLOR_FORCE")
+        .map(|var| var != "0")
+        .unwrap_or(false);
+    if no_color {
+        ColorMode::Never
+    } else if clicolor_force {
+        ColorMode::Always
+    } else {
+        let clicolor = std::env::var("CLICOLOR")
+            .map(|var| var != "0")
+            .unwrap_or(true);
+        if clicolor {
+            ColorMode::Auto
+        } else {
+            ColorMode::Never
+        }
+    }
+});
+
+/// The color output mode for a `ConsoleAppender`
+#[derive(Clone, Copy, Default)]
+pub enum ColorMode {
+    /// Print color only if the output is recognized as a console
+    #[default]
+    Auto,
+    /// Force color output
+    Always,
+    /// Never print color
+    Never,
+}
 
 /// An `encode::Write`r that outputs to a console.
 pub struct ConsoleWriter(imp::Writer);
@@ -88,7 +124,14 @@ mod imp {
     use std::{fmt, io};
 
     use crate::{
-        encode::{self, writer::ansi::AnsiWriter, Style},
+        encode::{
+            self,
+            writer::{
+                ansi::AnsiWriter,
+                console::{ColorMode, COLOR_MODE},
+            },
+            Style,
+        },
         priv_io::{StdWriter, StdWriterLock},
     };
 
@@ -96,19 +139,33 @@ mod imp {
 
     impl Writer {
         pub fn stdout() -> Option<Writer> {
-            if unsafe { libc::isatty(libc::STDOUT_FILENO) } != 1 {
-                return None;
+            let writer = || Writer(AnsiWriter(StdWriter::stdout()));
+            match *COLOR_MODE {
+                ColorMode::Auto => {
+                    if unsafe { libc::isatty(libc::STDOUT_FILENO) } != 1 {
+                        None
+                    } else {
+                        Some(writer())
+                    }
+                }
+                ColorMode::Always => Some(writer()),
+                ColorMode::Never => None,
             }
-
-            Some(Writer(AnsiWriter(StdWriter::stdout())))
         }
 
         pub fn stderr() -> Option<Writer> {
-            if unsafe { libc::isatty(libc::STDERR_FILENO) } != 1 {
-                return None;
+            let writer = || Writer(AnsiWriter(StdWriter::stderr()));
+            match *COLOR_MODE {
+                ColorMode::Auto => {
+                    if unsafe { libc::isatty(libc::STDERR_FILENO) } != 1 {
+                        None
+                    } else {
+                        Some(writer())
+                    }
+                }
+                ColorMode::Always => Some(writer()),
+                ColorMode::Never => None,
             }
-
-            Some(Writer(AnsiWriter(StdWriter::stderr())))
         }
 
         pub fn lock(&self) -> WriterLock {
@@ -180,7 +237,11 @@ mod imp {
     };
 
     use crate::{
-        encode::{self, Color, Style},
+        encode::{
+            self,
+            writer::console::{ColorMode, COLOR_MODE},
+            Color, Style,
+        },
         priv_io::{StdWriter, StdWriterLock},
     };
 
@@ -266,13 +327,18 @@ mod imp {
                     return None;
                 }
 
-                Some(Writer {
+                let writer = Writer {
                     console: RawConsole {
                         handle,
                         defaults: info.wAttributes,
                     },
                     inner: StdWriter::stdout(),
-                })
+                };
+
+                match *COLOR_MODE {
+                    ColorMode::Auto | ColorMode::Always => Some(writer),
+                    ColorMode::Never => None,
+                }
             }
         }
 
@@ -288,13 +354,18 @@ mod imp {
                     return None;
                 }
 
-                Some(Writer {
+                let writer = Writer {
                     console: RawConsole {
                         handle,
                         defaults: info.wAttributes,
                     },
-                    inner: StdWriter::stderr(),
-                })
+                    inner: StdWriter::stdout(),
+                };
+
+                match *COLOR_MODE {
+                    ColorMode::Auto | ColorMode::Always => Some(writer),
+                    ColorMode::Never => None,
+                }
             }
         }
 
diff --git a/tests/color_control.rs b/tests/color_control.rs
new file mode 100644
index 00000000..344f032f
--- /dev/null
+++ b/tests/color_control.rs
@@ -0,0 +1,24 @@
+use std::process::Command;
+
+fn execute_test(env_key: &str, env_val: &str) {
+    let mut child_proc = Command::new("cargo")
+        .args(&["run", "--example", "compile_time_config"])
+        .env(env_key, env_val)
+        .spawn()
+        .expect("Cargo command failed to start");
+
+    let ecode = child_proc.wait().expect("failed to wait on child");
+
+    assert!(ecode.success());
+}
+
+// Maintaining as a single test to avoid blocking calls to the package cache
+#[test]
+fn test_no_color() {
+    let keys = vec!["NO_COLOR", "CLICOLOR_FORCE", "CLICOLOR"];
+
+    for key in keys {
+        execute_test(key, "1");
+        execute_test(key, "0");
+    }
+}