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

colored/
control.rs
