use std::io::IsTerminal;
use std::sync::atomic::{AtomicU8, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColorChoice {
Auto,
Always,
Never,
}
const AUTO: u8 = 0;
const ALWAYS: u8 = 1;
const NEVER: u8 = 2;
static CHOICE: AtomicU8 = AtomicU8::new(AUTO);
pub fn set_color_choice(choice: ColorChoice) {
let v = match choice {
ColorChoice::Auto => AUTO,
ColorChoice::Always => ALWAYS,
ColorChoice::Never => NEVER,
};
CHOICE.store(v, Ordering::Relaxed);
}
pub fn color_choice() -> ColorChoice {
match CHOICE.load(Ordering::Relaxed) {
ALWAYS => ColorChoice::Always,
NEVER => ColorChoice::Never,
_ => ColorChoice::Auto,
}
}
pub fn color_enabled() -> bool {
match color_choice() {
ColorChoice::Always => true,
ColorChoice::Never => false,
ColorChoice::Auto => {
if env_flag_set("NO_COLOR") {
return false;
}
if env_flag_set("FORCE_COLOR") {
return true;
}
std::io::stderr().is_terminal()
}
}
}
fn env_flag_set(name: &str) -> bool {
std::env::var_os(name)
.map(|v| !v.is_empty())
.unwrap_or(false)
}
pub fn green(s: &str) -> String {
if color_enabled() {
format!("\x1b[32m{s}\x1b[0m")
} else {
s.to_string()
}
}
pub fn yellow(s: &str) -> String {
if color_enabled() {
format!("\x1b[33m{s}\x1b[0m")
} else {
s.to_string()
}
}
pub fn bold(s: &str) -> String {
if color_enabled() {
format!("\x1b[1m{s}\x1b[0m")
} else {
s.to_string()
}
}
pub fn dim(s: &str) -> String {
if color_enabled() {
format!("\x1b[2m{s}\x1b[0m")
} else {
s.to_string()
}
}
pub fn red(s: &str) -> String {
if color_enabled() {
format!("\x1b[31m{s}\x1b[0m")
} else {
s.to_string()
}
}
pub fn print_cli_error(msg: impl std::fmt::Display) {
eprintln!("{}: {msg}", red("error"));
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
static LOCK: Mutex<()> = Mutex::new(());
fn reset() {
set_color_choice(ColorChoice::Auto);
unsafe {
std::env::remove_var("NO_COLOR");
std::env::remove_var("FORCE_COLOR");
}
}
#[test]
fn explicit_always_forces_on() {
let _g = LOCK.lock().unwrap();
reset();
set_color_choice(ColorChoice::Always);
assert!(color_enabled());
reset();
}
#[test]
fn explicit_never_forces_off() {
let _g = LOCK.lock().unwrap();
reset();
set_color_choice(ColorChoice::Never);
assert!(!color_enabled());
reset();
}
#[test]
fn no_color_env_disables_in_auto_mode() {
let _g = LOCK.lock().unwrap();
reset();
unsafe {
std::env::set_var("NO_COLOR", "1");
}
assert!(!color_enabled());
reset();
}
#[test]
fn force_color_env_enables_in_auto_mode() {
let _g = LOCK.lock().unwrap();
reset();
unsafe {
std::env::set_var("FORCE_COLOR", "1");
}
assert!(color_enabled());
reset();
}
#[test]
fn no_color_beats_force_color_when_both_set() {
let _g = LOCK.lock().unwrap();
reset();
unsafe {
std::env::set_var("NO_COLOR", "1");
std::env::set_var("FORCE_COLOR", "1");
}
assert!(!color_enabled());
reset();
}
#[test]
fn explicit_override_beats_env_vars() {
let _g = LOCK.lock().unwrap();
reset();
unsafe {
std::env::set_var("NO_COLOR", "1");
}
set_color_choice(ColorChoice::Always);
assert!(color_enabled());
reset();
}
#[test]
fn empty_env_var_treated_as_unset() {
let _g = LOCK.lock().unwrap();
reset();
unsafe {
std::env::set_var("NO_COLOR", "");
}
assert_eq!(color_choice(), ColorChoice::Auto);
reset();
}
#[test]
fn green_yellow_bold_dim_empty_when_disabled() {
let _g = LOCK.lock().unwrap();
reset();
set_color_choice(ColorChoice::Never);
assert_eq!(green("x"), "x");
assert_eq!(yellow("x"), "x");
assert_eq!(bold("x"), "x");
assert_eq!(dim("x"), "x");
reset();
}
#[test]
fn green_wraps_with_ansi_when_enabled() {
let _g = LOCK.lock().unwrap();
reset();
set_color_choice(ColorChoice::Always);
assert_eq!(green("x"), "\x1b[32mx\x1b[0m");
reset();
}
}