1use std::io::IsTerminal;
9use std::sync::atomic::{AtomicU8, Ordering};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum ColorChoice {
15 Auto,
16 Always,
17 Never,
18}
19
20const AUTO: u8 = 0;
24const ALWAYS: u8 = 1;
25const NEVER: u8 = 2;
26
27static CHOICE: AtomicU8 = AtomicU8::new(AUTO);
28
29pub fn set_color_choice(choice: ColorChoice) {
33 let v = match choice {
34 ColorChoice::Auto => AUTO,
35 ColorChoice::Always => ALWAYS,
36 ColorChoice::Never => NEVER,
37 };
38 CHOICE.store(v, Ordering::Relaxed);
39}
40
41pub fn color_choice() -> ColorChoice {
43 match CHOICE.load(Ordering::Relaxed) {
44 ALWAYS => ColorChoice::Always,
45 NEVER => ColorChoice::Never,
46 _ => ColorChoice::Auto,
47 }
48}
49
50pub fn color_enabled() -> bool {
60 match color_choice() {
61 ColorChoice::Always => true,
62 ColorChoice::Never => false,
63 ColorChoice::Auto => {
64 if env_flag_set("NO_COLOR") {
65 return false;
66 }
67 if env_flag_set("FORCE_COLOR") {
68 return true;
69 }
70 std::io::stderr().is_terminal()
71 }
72 }
73}
74
75fn env_flag_set(name: &str) -> bool {
79 std::env::var_os(name)
80 .map(|v| !v.is_empty())
81 .unwrap_or(false)
82}
83
84pub fn green(s: &str) -> String {
87 if color_enabled() {
88 format!("\x1b[32m{s}\x1b[0m")
89 } else {
90 s.to_string()
91 }
92}
93
94pub fn yellow(s: &str) -> String {
95 if color_enabled() {
96 format!("\x1b[33m{s}\x1b[0m")
97 } else {
98 s.to_string()
99 }
100}
101
102pub fn bold(s: &str) -> String {
103 if color_enabled() {
104 format!("\x1b[1m{s}\x1b[0m")
105 } else {
106 s.to_string()
107 }
108}
109
110pub fn dim(s: &str) -> String {
111 if color_enabled() {
112 format!("\x1b[2m{s}\x1b[0m")
113 } else {
114 s.to_string()
115 }
116}
117
118pub fn red(s: &str) -> String {
119 if color_enabled() {
120 format!("\x1b[31m{s}\x1b[0m")
121 } else {
122 s.to_string()
123 }
124}
125
126pub fn print_cli_error(msg: impl std::fmt::Display) {
132 eprintln!("{}: {msg}", red("error"));
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use std::sync::Mutex;
139
140 static LOCK: Mutex<()> = Mutex::new(());
144
145 fn reset() {
146 set_color_choice(ColorChoice::Auto);
147 unsafe {
149 std::env::remove_var("NO_COLOR");
150 std::env::remove_var("FORCE_COLOR");
151 }
152 }
153
154 #[test]
155 fn explicit_always_forces_on() {
156 let _g = LOCK.lock().unwrap();
157 reset();
158 set_color_choice(ColorChoice::Always);
159 assert!(color_enabled());
160 reset();
161 }
162
163 #[test]
164 fn explicit_never_forces_off() {
165 let _g = LOCK.lock().unwrap();
166 reset();
167 set_color_choice(ColorChoice::Never);
168 assert!(!color_enabled());
169 reset();
170 }
171
172 #[test]
173 fn no_color_env_disables_in_auto_mode() {
174 let _g = LOCK.lock().unwrap();
175 reset();
176 unsafe {
177 std::env::set_var("NO_COLOR", "1");
178 }
179 assert!(!color_enabled());
180 reset();
181 }
182
183 #[test]
184 fn force_color_env_enables_in_auto_mode() {
185 let _g = LOCK.lock().unwrap();
186 reset();
187 unsafe {
188 std::env::set_var("FORCE_COLOR", "1");
189 }
190 assert!(color_enabled());
191 reset();
192 }
193
194 #[test]
195 fn no_color_beats_force_color_when_both_set() {
196 let _g = LOCK.lock().unwrap();
200 reset();
201 unsafe {
202 std::env::set_var("NO_COLOR", "1");
203 std::env::set_var("FORCE_COLOR", "1");
204 }
205 assert!(!color_enabled());
206 reset();
207 }
208
209 #[test]
210 fn explicit_override_beats_env_vars() {
211 let _g = LOCK.lock().unwrap();
212 reset();
213 unsafe {
214 std::env::set_var("NO_COLOR", "1");
215 }
216 set_color_choice(ColorChoice::Always);
217 assert!(color_enabled());
218 reset();
219 }
220
221 #[test]
222 fn empty_env_var_treated_as_unset() {
223 let _g = LOCK.lock().unwrap();
224 reset();
225 unsafe {
226 std::env::set_var("NO_COLOR", "");
227 }
228 assert_eq!(color_choice(), ColorChoice::Auto);
232 reset();
233 }
234
235 #[test]
236 fn green_yellow_bold_dim_empty_when_disabled() {
237 let _g = LOCK.lock().unwrap();
238 reset();
239 set_color_choice(ColorChoice::Never);
240 assert_eq!(green("x"), "x");
241 assert_eq!(yellow("x"), "x");
242 assert_eq!(bold("x"), "x");
243 assert_eq!(dim("x"), "x");
244 reset();
245 }
246
247 #[test]
248 fn green_wraps_with_ansi_when_enabled() {
249 let _g = LOCK.lock().unwrap();
250 reset();
251 set_color_choice(ColorChoice::Always);
252 assert_eq!(green("x"), "\x1b[32mx\x1b[0m");
253 reset();
254 }
255}