owo_colors/
overrides.rs

1use core::sync::atomic::{AtomicU8, Ordering};
2
3/// Set an override value for whether or not colors are supported using
4/// [`set_override`] while executing the closure provided.
5///
6/// Once the function has executed the value will be reset to the previous set
7/// (or unset) override.
8///
9/// This is especially useful in use-cases where one would like to temporarily
10/// override the supported color set, without impacting previous configurations.
11///
12/// ```
13/// # use owo_colors::{Stream, OwoColorize, set_override, unset_override, with_override};
14/// # use owo_colors::colors::Black;
15/// #
16/// set_override(false);
17/// assert_eq!("example".if_supports_color(Stream::Stdout, |value| value.bg::<Black>()).to_string(), "example");
18///
19/// with_override(true, || {
20///     assert_eq!("example".if_supports_color(Stream::Stdout, |value| value.bg::<Black>()).to_string(), "\x1b[40mexample\x1b[49m");
21/// });
22///
23/// assert_eq!("example".if_supports_color(Stream::Stdout, |value| value.bg::<Black>()).to_string(), "example");
24/// # unset_override() // make sure that other doc tests are not impacted
25/// ```
26#[cfg(feature = "supports-colors")]
27pub fn with_override<T, F: FnOnce() -> T>(enabled: bool, f: F) -> T {
28    let previous = OVERRIDE.inner();
29    OVERRIDE.set_force(enabled);
30
31    // Use a scope guard to ensure that if `f` panics, the override is still
32    // caught.
33    let _guard = ResetOverrideGuard { previous };
34
35    f()
36}
37
38struct ResetOverrideGuard {
39    previous: u8,
40}
41
42impl Drop for ResetOverrideGuard {
43    fn drop(&mut self) {
44        OVERRIDE.set_unchecked(self.previous);
45    }
46}
47
48/// Set an override value for whether or not colors are supported.
49///
50/// If `true` is passed,
51/// [`if_supports_color`](crate::OwoColorize::if_supports_color) will always act
52/// as if colors are supported.
53///
54/// If `false` is passed,
55/// [`if_supports_color`](crate::OwoColorize::if_supports_color) will always act
56/// as if colors are **not** supported.
57///
58/// This behavior can be disabled using [`unset_override`], allowing
59/// `owo-colors` to return to inferring if colors are supported.
60#[cfg(feature = "supports-colors")]
61pub fn set_override(enabled: bool) {
62    OVERRIDE.set_force(enabled);
63}
64
65/// Remove any override value for whether or not colors are supported. This
66/// means [`if_supports_color`](crate::OwoColorize::if_supports_color) will
67/// resume checking if the given terminal output ([`Stream`](crate::Stream))
68/// supports colors.
69///
70/// This override can be set using [`set_override`].
71#[cfg(feature = "supports-colors")]
72pub fn unset_override() {
73    OVERRIDE.unset();
74}
75
76pub(crate) static OVERRIDE: Override = Override::none();
77
78pub(crate) struct Override(AtomicU8);
79
80const FORCE_MASK: u8 = 0b10;
81const FORCE_ENABLE: u8 = 0b11;
82const FORCE_DISABLE: u8 = 0b10;
83const NO_FORCE: u8 = 0b00;
84
85impl Override {
86    const fn none() -> Self {
87        Self(AtomicU8::new(NO_FORCE))
88    }
89
90    fn inner(&self) -> u8 {
91        self.0.load(Ordering::SeqCst)
92    }
93
94    pub(crate) fn is_force_enabled_or_disabled(&self) -> (bool, bool) {
95        let inner = self.inner();
96
97        (inner == FORCE_ENABLE, inner == FORCE_DISABLE)
98    }
99
100    fn set_force(&self, enable: bool) {
101        self.set_unchecked(FORCE_MASK | (enable as u8));
102    }
103
104    fn unset(&self) {
105        self.set_unchecked(0);
106    }
107
108    fn set_unchecked(&self, value: u8) {
109        self.0.store(value, Ordering::SeqCst);
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn with_override_on_panic() {
119        set_override(false);
120
121        std::panic::catch_unwind(|| {
122            with_override(true, || {
123                assert_eq!(OVERRIDE.inner(), FORCE_ENABLE);
124                panic!("test");
125            });
126        })
127        .expect_err("test should panic");
128
129        assert_eq!(
130            OVERRIDE.inner(),
131            FORCE_DISABLE,
132            "override should have been reset"
133        );
134    }
135}
owo_colors/overrides.rs

owo_colors/
overrides.rs
