colored/
control.rs

1//! A couple of functions to enable and disable coloring.
2
3use std::default::Default;
4use std::env;
5use std::io::{self, IsTerminal};
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::LazyLock;
8
9/// Sets a flag to the console to use a virtual terminal environment.
10///
11/// This is primarily used for Windows 10 environments which will not correctly colorize
12/// the outputs based on ANSI escape codes.
13///
14/// The returned `Result` is _always_ `Ok(())`, the return type was kept to ensure backwards
15/// compatibility.
16///
17/// # Notes
18/// > Only available to `Windows` build targets.
19///
20/// # Example
21/// ```rust
22/// use colored::*;
23/// control::set_virtual_terminal(false).unwrap();
24/// println!("{}", "bright cyan".bright_cyan());    // will print '[96mbright cyan[0m' on windows 10
25///
26/// control::set_virtual_terminal(true).unwrap();
27/// println!("{}", "bright cyan".bright_cyan());    // will print correctly
28/// ```
29#[allow(clippy::result_unit_err)]
30#[cfg(windows)]
31pub fn set_virtual_terminal(use_virtual: bool) -> Result<(), ()> {
32    use windows_sys::Win32::System::Console::{
33        GetConsoleMode, GetStdHandle, SetConsoleMode, ENABLE_VIRTUAL_TERMINAL_PROCESSING,
34        STD_OUTPUT_HANDLE,
35    };
36
37    unsafe {
38        let handle = GetStdHandle(STD_OUTPUT_HANDLE);
39        let mut original_mode = 0;
40        GetConsoleMode(handle, &mut original_mode);
41
42        let enabled = original_mode & ENABLE_VIRTUAL_TERMINAL_PROCESSING
43            == ENABLE_VIRTUAL_TERMINAL_PROCESSING;
44
45        match (use_virtual, enabled) {
46            // not enabled, should be enabled
47            (true, false) => {
48                SetConsoleMode(handle, ENABLE_VIRTUAL_TERMINAL_PROCESSING | original_mode)
49            }
50            // already enabled, should be disabled
51            (false, true) => {
52                SetConsoleMode(handle, ENABLE_VIRTUAL_TERMINAL_PROCESSING ^ original_mode)
53            }
54            _ => 0,
55        };
56    }
57
58    Ok(())
59}
60
61/// A flag for whether coloring should occur.
62pub struct ShouldColorize {
63    clicolor: bool,
64    clicolor_force: Option<bool>,
65    // XXX we can't use Option<Atomic> because we can't use &mut references to ShouldColorize
66    has_manual_override: AtomicBool,
67    manual_override: AtomicBool,
68}
69
70/// Use this to force colored to ignore the environment and always/never colorize
71/// See example/control.rs
72pub fn set_override(override_colorize: bool) {
73    SHOULD_COLORIZE.set_override(override_colorize);
74}
75
76/// Remove the manual override and let the environment decide if it's ok to colorize
77/// See example/control.rs
78pub fn unset_override() {
79    SHOULD_COLORIZE.unset_override();
80}
81
82/// The persistent [`ShouldColorize`].
83pub static SHOULD_COLORIZE: LazyLock<ShouldColorize> = LazyLock::new(ShouldColorize::from_env);
84
85impl Default for ShouldColorize {
86    fn default() -> ShouldColorize {
87        ShouldColorize {
88            clicolor: true,
89            clicolor_force: None,
90            has_manual_override: AtomicBool::new(false),
91            manual_override: AtomicBool::new(false),
92        }
93    }
94}
95
96impl ShouldColorize {
97    /// Reads environment variables and checks if output is a tty to determine
98    /// whether colorization should be used or not.
99    /// `CLICOLOR_FORCE` takes highest priority, followed by `NO_COLOR`,
100    /// followed by `CLICOLOR` combined with tty check.
101    pub fn from_env() -> Self {
102        ShouldColorize {
103            clicolor: ShouldColorize::normalize_env(env::var("CLICOLOR")).unwrap_or(true)
104                && io::stdout().is_terminal(),
105            clicolor_force: ShouldColorize::resolve_clicolor_force(
106                env::var("NO_COLOR"),
107                env::var("CLICOLOR_FORCE"),
108            ),
109            ..ShouldColorize::default()
110        }
111    }
112
113    /// Returns if the current coloring is expected.
114    pub fn should_colorize(&self) -> bool {
115        if self.has_manual_override.load(Ordering::Relaxed) {
116            return self.manual_override.load(Ordering::Relaxed);
117        }
118
119        if let Some(forced_value) = self.clicolor_force {
120            return forced_value;
121        }
122
123        self.clicolor
124    }
125
126    /// Use this to force colored to ignore the environment and always/never colorize
127    pub fn set_override(&self, override_colorize: bool) {
128        self.has_manual_override.store(true, Ordering::Relaxed);
129        self.manual_override
130            .store(override_colorize, Ordering::Relaxed);
131    }
132
133    /// Remove the manual override and let the environment decide if it's ok to colorize
134    pub fn unset_override(&self) {
135        self.has_manual_override.store(false, Ordering::Relaxed);
136    }
137
138    /* private */
139
140    fn normalize_env(env_res: Result<String, env::VarError>) -> Option<bool> {
141        match env_res {
142            Ok(string) => Some(string != "0"),
143            Err(_) => None,
144        }
145    }
146
147    fn resolve_clicolor_force(
148        no_color: Result<String, env::VarError>,
149        clicolor_force: Result<String, env::VarError>,
150    ) -> Option<bool> {
151        if ShouldColorize::normalize_env(clicolor_force) == Some(true) {
152            Some(true)
153        } else if ShouldColorize::normalize_env(no_color).is_some() {
154            Some(false)
155        } else {
156            None
157        }
158    }
159}
160
161#[cfg(test)]
162mod specs {
163    use super::*;
164    use rspec;
165    use std::env;
166
167    #[test]
168    fn clicolor_behavior() {
169        rspec::run(&rspec::describe("ShouldColorize", (), |ctx| {
170            ctx.specify("::normalize_env", |ctx| {
171                ctx.it("should return None if error", |_| {
172                    assert_eq!(
173                        None,
174                        ShouldColorize::normalize_env(Err(env::VarError::NotPresent))
175                    );
176                    assert_eq!(
177                        None,
178                        ShouldColorize::normalize_env(Err(env::VarError::NotUnicode("".into())))
179                    );
180                });
181
182                ctx.it("should return Some(true) if != 0", |_| {
183                    Some(true) == ShouldColorize::normalize_env(Ok(String::from("1")))
184                });
185
186                ctx.it("should return Some(false) if == 0", |_| {
187                    Some(false) == ShouldColorize::normalize_env(Ok(String::from("0")))
188                });
189            });
190
191            ctx.specify("::resolve_clicolor_force", |ctx| {
192                ctx.it(
193                    "should return None if NO_COLOR is not set and CLICOLOR_FORCE is not set or set to 0",
194                    |_| {
195                        assert_eq!(
196                            None,
197                            ShouldColorize::resolve_clicolor_force(
198                                Err(env::VarError::NotPresent),
199                                Err(env::VarError::NotPresent)
200                            )
201                        );
202                        assert_eq!(
203                            None,
204                            ShouldColorize::resolve_clicolor_force(
205                                Err(env::VarError::NotPresent),
206                                Ok(String::from("0")),
207                            )
208                        );
209                    },
210                );
211
212                ctx.it(
213                    "should return Some(false) if NO_COLOR is set and CLICOLOR_FORCE is not enabled",
214                    |_| {
215                        assert_eq!(
216                            Some(false),
217                            ShouldColorize::resolve_clicolor_force(
218                                Ok(String::from("0")),
219                                Err(env::VarError::NotPresent)
220                            )
221                        );
222                        assert_eq!(
223                            Some(false),
224                            ShouldColorize::resolve_clicolor_force(
225                                Ok(String::from("1")),
226                                Err(env::VarError::NotPresent)
227                            )
228                        );
229                        assert_eq!(
230                            Some(false),
231                            ShouldColorize::resolve_clicolor_force(
232                                Ok(String::from("1")),
233                                Ok(String::from("0")),
234                            )
235                        );
236                    },
237                );
238
239                ctx.it(
240                    "should prioritize CLICOLOR_FORCE over NO_COLOR if CLICOLOR_FORCE is set to non-zero value",
241                    |_| {
242                        assert_eq!(
243                            Some(true),
244                            ShouldColorize::resolve_clicolor_force(
245                                Ok(String::from("1")),
246                                Ok(String::from("1")),
247                            )
248                        );
249                        assert_eq!(
250                            Some(false),
251                            ShouldColorize::resolve_clicolor_force(
252                                Ok(String::from("1")),
253                                Ok(String::from("0")),
254                            )
255                        );
256                        assert_eq!(
257                            Some(true),
258                            ShouldColorize::resolve_clicolor_force(
259                                Err(env::VarError::NotPresent),
260                                Ok(String::from("1")),
261                            )
262                        );
263                    },
264                );
265            });
266
267            ctx.specify("constructors", |ctx| {
268                ctx.it("should have a default constructor", |_| {
269                    ShouldColorize::default();
270                });
271
272                ctx.it("should have an environment constructor", |_| {
273                    ShouldColorize::from_env();
274                });
275            });
276
277            ctx.specify("when only changing clicolors", |ctx| {
278                ctx.it("clicolor == false means no colors", |_| {
279                    let colorize_control = ShouldColorize {
280                        clicolor: false,
281                        ..ShouldColorize::default()
282                    };
283                    !colorize_control.should_colorize()
284                });
285
286                ctx.it("clicolor == true means colors !", |_| {
287                    let colorize_control = ShouldColorize {
288                        clicolor: true,
289                        ..ShouldColorize::default()
290                    };
291                    colorize_control.should_colorize()
292                });
293
294                ctx.it("unset clicolors implies true", |_| {
295                    ShouldColorize::default().should_colorize()
296                });
297            });
298
299            ctx.specify("when using clicolor_force", |ctx| {
300                ctx.it(
301                    "clicolor_force should force to true no matter clicolor",
302                    |_| {
303                        let colorize_control = ShouldColorize {
304                            clicolor: false,
305                            clicolor_force: Some(true),
306                            ..ShouldColorize::default()
307                        };
308
309                        colorize_control.should_colorize()
310                    },
311                );
312
313                ctx.it(
314                    "clicolor_force should force to false no matter clicolor",
315                    |_| {
316                        let colorize_control = ShouldColorize {
317                            clicolor: true,
318                            clicolor_force: Some(false),
319                            ..ShouldColorize::default()
320                        };
321
322                        !colorize_control.should_colorize()
323                    },
324                );
325            });
326
327            ctx.specify("using a manual override", |ctx| {
328                ctx.it("shoud colorize if manual_override is true, but clicolor is false and clicolor_force also false", |_| {
329                    let colorize_control = ShouldColorize {
330                        clicolor: false,
331                        clicolor_force: None,
332                        has_manual_override: AtomicBool::new(true),
333                        manual_override: AtomicBool::new(true),
334                    };
335
336                    colorize_control.should_colorize();
337                });
338
339                ctx.it("should not colorize if manual_override is false, but clicolor is true or clicolor_force is true", |_| {
340                    let colorize_control = ShouldColorize {
341                        clicolor: true,
342                        clicolor_force: Some(true),
343                        has_manual_override: AtomicBool::new(true),
344                        manual_override: AtomicBool::new(false),
345                    };
346
347                    !colorize_control.should_colorize()
348                });
349            });
350
351            ctx.specify("::set_override", |ctx| {
352                ctx.it("should exists", |_| {
353                    let colorize_control = ShouldColorize::default();
354                    colorize_control.set_override(true);
355                });
356
357                ctx.it("set the manual_override property", |_| {
358                    let colorize_control = ShouldColorize::default();
359                    colorize_control.set_override(true);
360                    {
361                        assert!(colorize_control.has_manual_override.load(Ordering::Relaxed));
362                        let val = colorize_control.manual_override.load(Ordering::Relaxed);
363                        assert!(val);
364                    }
365                    colorize_control.set_override(false);
366                    {
367                        assert!(colorize_control.has_manual_override.load(Ordering::Relaxed));
368                        let val = colorize_control.manual_override.load(Ordering::Relaxed);
369                        assert!(!val);
370                    }
371                });
372            });
373
374            ctx.specify("::unset_override", |ctx| {
375                ctx.it("should exists", |_| {
376                    let colorize_control = ShouldColorize::default();
377                    colorize_control.unset_override();
378                });
379
380                ctx.it("unset the manual_override property", |_| {
381                    let colorize_control = ShouldColorize::default();
382                    colorize_control.set_override(true);
383                    colorize_control.unset_override();
384                    assert!(!colorize_control.has_manual_override.load(Ordering::Relaxed));
385                });
386            });
387        }));
388    }
389}
colored/control.rs

colored/
control.rs
