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 'bright cyan' 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}
391

Provided by KDAB

Privacy Policy
Learn Rust with the experts
Find out more