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)]
141pub(crate) static TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146 use super::TEST_ENV_LOCK as LOCK;
147
148 fn reset() {
149 set_color_choice(ColorChoice::Auto);
150 unsafe {
152 std::env::remove_var("NO_COLOR");
153 std::env::remove_var("FORCE_COLOR");
154 }
155 }
156
157 #[test]
158 fn explicit_always_forces_on() {
159 let _g = LOCK.lock().unwrap();
160 reset();
161 set_color_choice(ColorChoice::Always);
162 assert!(color_enabled());
163 reset();
164 }
165
166 #[test]
167 fn explicit_never_forces_off() {
168 let _g = LOCK.lock().unwrap();
169 reset();
170 set_color_choice(ColorChoice::Never);
171 assert!(!color_enabled());
172 reset();
173 }
174
175 #[test]
176 fn no_color_env_disables_in_auto_mode() {
177 let _g = LOCK.lock().unwrap();
178 reset();
179 unsafe {
180 std::env::set_var("NO_COLOR", "1");
181 }
182 assert!(!color_enabled());
183 reset();
184 }
185
186 #[test]
187 fn force_color_env_enables_in_auto_mode() {
188 let _g = LOCK.lock().unwrap();
189 reset();
190 unsafe {
191 std::env::set_var("FORCE_COLOR", "1");
192 }
193 assert!(color_enabled());
194 reset();
195 }
196
197 #[test]
198 fn no_color_beats_force_color_when_both_set() {
199 let _g = LOCK.lock().unwrap();
203 reset();
204 unsafe {
205 std::env::set_var("NO_COLOR", "1");
206 std::env::set_var("FORCE_COLOR", "1");
207 }
208 assert!(!color_enabled());
209 reset();
210 }
211
212 #[test]
213 fn explicit_override_beats_env_vars() {
214 let _g = LOCK.lock().unwrap();
215 reset();
216 unsafe {
217 std::env::set_var("NO_COLOR", "1");
218 }
219 set_color_choice(ColorChoice::Always);
220 assert!(color_enabled());
221 reset();
222 }
223
224 #[test]
225 fn empty_env_var_treated_as_unset() {
226 let _g = LOCK.lock().unwrap();
227 reset();
228 unsafe {
229 std::env::set_var("NO_COLOR", "");
230 }
231 assert_eq!(color_choice(), ColorChoice::Auto);
235 reset();
236 }
237
238 #[test]
239 fn green_yellow_bold_dim_empty_when_disabled() {
240 let _g = LOCK.lock().unwrap();
241 reset();
242 set_color_choice(ColorChoice::Never);
243 assert_eq!(green("x"), "x");
244 assert_eq!(yellow("x"), "x");
245 assert_eq!(bold("x"), "x");
246 assert_eq!(dim("x"), "x");
247 reset();
248 }
249
250 #[test]
251 fn green_wraps_with_ansi_when_enabled() {
252 let _g = LOCK.lock().unwrap();
253 reset();
254 set_color_choice(ColorChoice::Always);
255 assert_eq!(green("x"), "\x1b[32mx\x1b[0m");
256 reset();
257 }
258}