1//! A couple of functions to enable and disable coloring.
2
3use is_terminal::IsTerminal;
4use std::default::Default;
5use std::env;
6use std::io;
7use std::sync::atomic::{AtomicBool, Ordering};
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 'bright cyan' 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
82lazy_static! {
83/// The persistent [`ShouldColorize`].
84 pub static ref SHOULD_COLORIZE: ShouldColorize = ShouldColorize::from_env();
85}
86
87impl Default for ShouldColorize {
88 fn default() -> ShouldColorize {
89 ShouldColorize {
90 clicolor: true,
91 clicolor_force: None,
92 has_manual_override: AtomicBool::new(false),
93 manual_override: AtomicBool::new(false),
94 }
95 }
96}
97
98impl ShouldColorize {
99 /// Reads environment variables and checks if output is a tty to determine
100 /// whether colorization should be used or not.
101 /// `CLICOLOR_FORCE` takes highest priority, followed by `NO_COLOR`,
102 /// followed by `CLICOLOR` combined with tty check.
103 pub fn from_env() -> Self {
104 ShouldColorize {
105 clicolor: ShouldColorize::normalize_env(env::var("CLICOLOR")).unwrap_or(true)
106 && io::stdout().is_terminal(),
107 clicolor_force: ShouldColorize::resolve_clicolor_force(
108 env::var("NO_COLOR"),
109 env::var("CLICOLOR_FORCE"),
110 ),
111 ..ShouldColorize::default()
112 }
113 }
114
115 /// Returns if the current coloring is expected.
116 pub fn should_colorize(&self) -> bool {
117 if self.has_manual_override.load(Ordering::Relaxed) {
118 return self.manual_override.load(Ordering::Relaxed);
119 }
120
121 if let Some(forced_value) = self.clicolor_force {
122 return forced_value;
123 }
124
125 self.clicolor
126 }
127
128 /// Use this to force colored to ignore the environment and always/never colorize
129 pub fn set_override(&self, override_colorize: bool) {
130 self.has_manual_override.store(true, Ordering::Relaxed);
131 self.manual_override
132 .store(override_colorize, Ordering::Relaxed);
133 }
134
135 /// Remove the manual override and let the environment decide if it's ok to colorize
136 pub fn unset_override(&self) {
137 self.has_manual_override.store(false, Ordering::Relaxed);
138 }
139
140 /* private */
141
142 fn normalize_env(env_res: Result<String, env::VarError>) -> Option<bool> {
143 match env_res {
144 Ok(string) => Some(string != "0"),
145 Err(_) => None,
146 }
147 }
148
149 fn resolve_clicolor_force(
150 no_color: Result<String, env::VarError>,
151 clicolor_force: Result<String, env::VarError>,
152 ) -> Option<bool> {
153 if ShouldColorize::normalize_env(clicolor_force) == Some(true) {
154 Some(true)
155 } else if ShouldColorize::normalize_env(no_color).is_some() {
156 Some(false)
157 } else {
158 None
159 }
160 }
161}
162
163#[cfg(test)]
164mod specs {
165 use super::*;
166 use rspec;
167 use rspec::context::*;
168 use std::env;
169
170 #[test]
171 fn clicolor_behavior() {
172 use std::io;
173
174 let stdout = &mut io::stdout();
175 let mut formatter = rspec::formatter::Simple::new(stdout);
176 let mut runner = describe("ShouldColorize", |ctx| {
177 ctx.describe("::normalize_env", |ctx| {
178 ctx.it("should return None if error", || {
179 assert_eq!(
180 None,
181 ShouldColorize::normalize_env(Err(env::VarError::NotPresent))
182 );
183 assert_eq!(
184 None,
185 ShouldColorize::normalize_env(Err(env::VarError::NotUnicode("".into())))
186 )
187 });
188
189 ctx.it("should return Some(true) if != 0", || {
190 Some(true) == ShouldColorize::normalize_env(Ok(String::from("1")))
191 });
192
193 ctx.it("should return Some(false) if == 0", || {
194 Some(false) == ShouldColorize::normalize_env(Ok(String::from("0")))
195 });
196 });
197
198 ctx.describe("::resolve_clicolor_force", |ctx| {
199 ctx.it(
200 "should return None if NO_COLOR is not set and CLICOLOR_FORCE is not set or set to 0",
201 || {
202 assert_eq!(
203 None,
204 ShouldColorize::resolve_clicolor_force(
205 Err(env::VarError::NotPresent),
206 Err(env::VarError::NotPresent)
207 )
208 );
209 assert_eq!(
210 None,
211 ShouldColorize::resolve_clicolor_force(
212 Err(env::VarError::NotPresent),
213 Ok(String::from("0")),
214 )
215 );
216 },
217 );
218
219 ctx.it(
220 "should return Some(false) if NO_COLOR is set and CLICOLOR_FORCE is not enabled",
221 || {
222 assert_eq!(
223 Some(false),
224 ShouldColorize::resolve_clicolor_force(
225 Ok(String::from("0")),
226 Err(env::VarError::NotPresent)
227 )
228 );
229 assert_eq!(
230 Some(false),
231 ShouldColorize::resolve_clicolor_force(
232 Ok(String::from("1")),
233 Err(env::VarError::NotPresent)
234 )
235 );
236 assert_eq!(
237 Some(false),
238 ShouldColorize::resolve_clicolor_force(
239 Ok(String::from("1")),
240 Ok(String::from("0")),
241 )
242 );
243 },
244 );
245
246 ctx.it(
247 "should prioritize CLICOLOR_FORCE over NO_COLOR if CLICOLOR_FORCE is set to non-zero value",
248 || {
249 assert_eq!(
250 Some(true),
251 ShouldColorize::resolve_clicolor_force(
252 Ok(String::from("1")),
253 Ok(String::from("1")),
254 )
255 );
256 assert_eq!(
257 Some(false),
258 ShouldColorize::resolve_clicolor_force(
259 Ok(String::from("1")),
260 Ok(String::from("0")),
261 )
262 );
263 assert_eq!(
264 Some(true),
265 ShouldColorize::resolve_clicolor_force(
266 Err(env::VarError::NotPresent),
267 Ok(String::from("1")),
268 )
269 );
270 },
271 );
272 });
273
274 ctx.describe("constructors", |ctx| {
275 ctx.it("should have a default constructor", || {
276 ShouldColorize::default();
277 });
278
279 ctx.it("should have an environment constructor", || {
280 ShouldColorize::from_env();
281 });
282 });
283
284 ctx.describe("when only changing clicolors", |ctx| {
285 ctx.it("clicolor == false means no colors", || {
286 let colorize_control = ShouldColorize {
287 clicolor: false,
288 ..ShouldColorize::default()
289 };
290 false == colorize_control.should_colorize()
291 });
292
293 ctx.it("clicolor == true means colors !", || {
294 let colorize_control = ShouldColorize {
295 clicolor: true,
296 ..ShouldColorize::default()
297 };
298 true == colorize_control.should_colorize()
299 });
300
301 ctx.it("unset clicolors implies true", || {
302 true == ShouldColorize::default().should_colorize()
303 });
304 });
305
306 ctx.describe("when using clicolor_force", |ctx| {
307 ctx.it(
308 "clicolor_force should force to true no matter clicolor",
309 || {
310 let colorize_control = ShouldColorize {
311 clicolor: false,
312 clicolor_force: Some(true),
313 ..ShouldColorize::default()
314 };
315
316 true == colorize_control.should_colorize()
317 },
318 );
319
320 ctx.it(
321 "clicolor_force should force to false no matter clicolor",
322 || {
323 let colorize_control = ShouldColorize {
324 clicolor: true,
325 clicolor_force: Some(false),
326 ..ShouldColorize::default()
327 };
328
329 false == colorize_control.should_colorize()
330 },
331 );
332 });
333
334 ctx.describe("using a manual override", |ctx| {
335 ctx.it("shoud colorize if manual_override is true, but clicolor is false and clicolor_force also false", || {
336 let colorize_control = ShouldColorize {
337 clicolor: false,
338 clicolor_force: None,
339 has_manual_override: AtomicBool::new(true),
340 manual_override: AtomicBool::new(true),
341 .. ShouldColorize::default()
342 };
343
344 true == colorize_control.should_colorize()
345 });
346
347 ctx.it("should not colorize if manual_override is false, but clicolor is true or clicolor_force is true", || {
348 let colorize_control = ShouldColorize {
349 clicolor: true,
350 clicolor_force: Some(true),
351 has_manual_override: AtomicBool::new(true),
352 manual_override: AtomicBool::new(false),
353 .. ShouldColorize::default()
354 };
355
356 false == colorize_control.should_colorize()
357 })
358 });
359
360 ctx.describe("::set_override", |ctx| {
361 ctx.it("should exists", || {
362 let colorize_control = ShouldColorize::default();
363 colorize_control.set_override(true);
364 });
365
366 ctx.it("set the manual_override property", || {
367 let colorize_control = ShouldColorize::default();
368 colorize_control.set_override(true);
369 {
370 assert_eq!(
371 true,
372 colorize_control.has_manual_override.load(Ordering::Relaxed)
373 );
374 let val = colorize_control.manual_override.load(Ordering::Relaxed);
375 assert_eq!(true, val);
376 }
377 colorize_control.set_override(false);
378 {
379 assert_eq!(
380 true,
381 colorize_control.has_manual_override.load(Ordering::Relaxed)
382 );
383 let val = colorize_control.manual_override.load(Ordering::Relaxed);
384 assert_eq!(false, val);
385 }
386 });
387 });
388
389 ctx.describe("::unset_override", |ctx| {
390 ctx.it("should exists", || {
391 let colorize_control = ShouldColorize::default();
392 colorize_control.unset_override();
393 });
394
395 ctx.it("unset the manual_override property", || {
396 let colorize_control = ShouldColorize::default();
397 colorize_control.set_override(true);
398 colorize_control.unset_override();
399 assert_eq!(
400 false,
401 colorize_control.has_manual_override.load(Ordering::Relaxed)
402 );
403 });
404 });
405 });
406 runner.add_event_handler(&mut formatter);
407 runner.run().unwrap();
408 }
409}
410