fancy_tree/color/
choice.rs

1//! Module for color choices that control what colors actually get displayed.
2use super::Color;
3use clap::ValueEnum;
4use mlua::{FromLua, Lua};
5use owo_colors::{
6    AnsiColors::{
7        self, Black, Blue, BrightBlack, BrightBlue, BrightCyan, BrightGreen, BrightMagenta,
8        BrightRed, BrightWhite, BrightYellow, Cyan, Green, Magenta, Red, White, Yellow,
9    },
10    DynColors, OwoColorize,
11    Stream::Stdout,
12};
13use std::fmt::Display;
14use std::io::{self, Write};
15
16/// Supports users choosing the colors they would like to display.
17#[derive(Debug, ValueEnum, Clone, Copy)]
18pub enum ColorChoice {
19    /// Let the application decide.
20    ///
21    /// *This checks if the `Stdout` stream supports colors.*
22    Auto,
23    /// Show all colors.
24    On,
25    /// Show only the 16 ANSI colors.
26    Ansi,
27    /// Don't show any colors.
28    Off,
29}
30
31impl ColorChoice {
32    /// Should colors support be automatically detected?
33    #[inline]
34    pub fn is_auto(&self) -> bool {
35        matches!(self, Self::Auto)
36    }
37
38    /// Should colors be on/enabled?
39    #[inline]
40    pub fn is_on(&self) -> bool {
41        matches!(self, Self::On)
42    }
43
44    /// Should colors be enabled but limited to the 16 ANSI colors?
45    #[inline]
46    pub fn is_ansi(&self) -> bool {
47        matches!(self, Self::Ansi)
48    }
49
50    /// Should colors be off/disabled?
51    #[inline]
52    pub fn is_off(&self) -> bool {
53        matches!(self, Self::Off)
54    }
55
56    /// Writes a colorized display value to the writer.
57    pub fn write_to<W, D>(
58        &self,
59        writer: &mut W,
60        display: D,
61        fg: Option<Color>,
62        bg: Option<Color>,
63    ) -> io::Result<()>
64    where
65        W: Write,
66        D: Display + OwoColorize,
67    {
68        match (self, fg, bg) {
69            // NOTE These variants must be on top
70            (Self::Off, _, _) | (_, None, None) => Self::off_write_to(writer, display),
71            (Self::Auto, fg, bg) => Self::auto_write_to(writer, display, fg, bg),
72            (Self::On, fg, bg) => Self::on_write_to(writer, display, fg, bg),
73            (Self::Ansi, fg, bg) => Self::ansi_write_to(writer, display, fg, bg),
74        }
75    }
76
77    /// Writes the display with color support detected.
78    fn auto_write_to<W, D, Fg, Bg>(
79        writer: &mut W,
80        display: D,
81        fg: Option<Fg>,
82        bg: Option<Bg>,
83    ) -> io::Result<()>
84    where
85        W: Write,
86        D: Display + OwoColorize,
87        DynColors: From<Fg>,
88        DynColors: From<Bg>,
89    {
90        let fg = fg.map(DynColors::from);
91        let bg = bg.map(DynColors::from);
92
93        // HACK This assumes that the writer is always Stdout, which might not be best
94        //      if we ever support other writers (Stderr, file, etc.).
95        match (fg, bg) {
96            (None, None) => unreachable!("Should use the off writer"),
97            (Some(fg), None) => write!(
98                writer,
99                "{}",
100                display.if_supports_color(Stdout, |display| display.color(fg))
101            ),
102            (None, Some(bg)) => write!(
103                writer,
104                "{}",
105                display.if_supports_color(Stdout, |display| display.on_color(bg))
106            ),
107            (Some(fg), Some(bg)) => write!(
108                writer,
109                "{}",
110                display.if_supports_color(Stdout, |display| display.color(fg).on_color(bg))
111            ),
112        }
113    }
114
115    /// Writes the display with no colorization.
116    #[inline]
117    fn off_write_to<W, D>(writer: &mut W, display: D) -> io::Result<()>
118    where
119        W: Write,
120        D: Display,
121    {
122        write!(writer, "{display}")
123    }
124
125    /// Writes the display with colorization on.
126    fn on_write_to<W, D, Fg, Bg>(
127        writer: &mut W,
128        display: D,
129        fg: Option<Fg>,
130        bg: Option<Bg>,
131    ) -> io::Result<()>
132    where
133        W: Write,
134        D: Display + OwoColorize,
135        DynColors: From<Fg>,
136        DynColors: From<Bg>,
137    {
138        let fg = fg.map(DynColors::from);
139        let bg = bg.map(DynColors::from);
140        match (fg, bg) {
141            (None, None) => unreachable!("Should use the off writer"),
142            (Some(fg), None) => write!(writer, "{}", display.color(fg)),
143            (None, Some(bg)) => write!(writer, "{}", display.on_color(bg)),
144            (Some(fg), Some(bg)) => write!(writer, "{}", display.color(fg).on_color(bg)),
145        }
146    }
147
148    /// Writes the display with colorization set to ANSI.
149    fn ansi_write_to<W, D>(
150        writer: &mut W,
151        display: D,
152        fg: Option<Color>,
153        bg: Option<Color>,
154    ) -> io::Result<()>
155    where
156        W: Write,
157        D: Display + OwoColorize,
158    {
159        let convert = |color: Color| DynColors::Ansi(Self::color_to_ansi(color));
160        let fg = fg.map(convert);
161        let bg = bg.map(convert);
162        Self::on_write_to(writer, display, fg, bg)
163    }
164
165    /// Converts the [`Color`] to ANSI.
166    fn color_to_ansi(color: Color) -> AnsiColors {
167        match color {
168            Color::Ansi(ansi) => ansi,
169            Color::Rgb(r, g, b) => Self::ansi_from_rgb(r, g, b),
170        }
171    }
172
173    /// Tries to get the closest ANSI color from RGB values.
174    fn ansi_from_rgb(r: u8, g: u8, b: u8) -> AnsiColors {
175        /// Stores colors to be indexed into by a 3-bit union of the RGB values.
176        const COLOR_INDEX: [AnsiColors; 16] = [
177            Black,
178            Red,
179            Green,
180            Yellow,
181            Blue,
182            Magenta,
183            Cyan,
184            White,
185            BrightBlack,
186            BrightRed,
187            BrightGreen,
188            BrightYellow,
189            BrightBlue,
190            BrightMagenta,
191            BrightCyan,
192            BrightWhite,
193        ];
194        /// Converts a color channel into a single bit at the given index.
195        #[inline]
196        const fn channel_bit(channel: u8, index: u8) -> u8 {
197            debug_assert!(index < 3);
198            (channel >> 7) << index
199        }
200
201        let brightness_index: usize = if Self::rgb_is_bright(r, g, b) { 8 } else { 0 };
202        let color_index = usize::from(channel_bit(r, 0) | channel_bit(g, 1) | channel_bit(b, 2));
203        debug_assert!(color_index <= 0b111);
204        let index = brightness_index + color_index;
205        debug_assert!(index < COLOR_INDEX.len());
206        COLOR_INDEX[index]
207    }
208
209    /// Detects if an RGB color is bright.
210    const fn rgb_is_bright(r: u8, g: u8, b: u8) -> bool {
211        /// Threshold between colors and their bright variants. If at least one color
212        /// channel is above this value, we assume the color is bright.
213        const BRIGHT_THRESHOLD: u8 = 0b1100_0000;
214
215        (r | g | b) >= BRIGHT_THRESHOLD || Self::rgb_is_bright_black(r, g, b)
216    }
217
218    /// Detects if an RGB color is bright black, which is a special case since bright black
219    /// is still dark compared to other colors.
220    const fn rgb_is_bright_black(r: u8, g: u8, b: u8) -> bool {
221        const BRIGHT_BLACK_MIN: u8 = 0b0100_0000;
222        const BRIGHT_BLACK_MAX: u8 = 0b1000_0000;
223        #[inline]
224        const fn within_limits(channel: u8) -> bool {
225            BRIGHT_BLACK_MIN <= channel && channel < BRIGHT_BLACK_MAX
226        }
227
228        within_limits(r) && within_limits(g) && within_limits(b)
229    }
230}
231
232impl Default for ColorChoice {
233    #[inline]
234    /// The auto variant.
235    fn default() -> Self {
236        Self::Auto
237    }
238}
239
240impl FromLua for ColorChoice {
241    fn from_lua(value: mlua::Value, _lua: &Lua) -> mlua::Result<Self> {
242        const VALID_VALUES: [&str; 4] = ["auto", "on", "off", "ansi"];
243        let type_name = value.type_name();
244        let make_conversion_error = || mlua::Error::FromLuaConversionError {
245            from: type_name,
246            to: String::from("ColorChoice"),
247            message: Some(format!("Must be one of {VALID_VALUES:?} or nil")),
248        };
249        let color_choice = value
250            .as_string()
251            .ok_or_else(make_conversion_error)?
252            .to_string_lossy();
253        let color_choice = color_choice.as_str();
254        let color_choice = match color_choice {
255            "auto" => Self::Auto,
256            "on" => Self::On,
257            "off" => Self::Off,
258            "ansi" => Self::Ansi,
259            _ => return Err(make_conversion_error()),
260        };
261        Ok(color_choice)
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use rstest::rstest;
269
270    #[rstest]
271    #[case::bright_black(85, 85, 85, BrightBlack)]
272    #[case::black(0, 0b10, 0b11, Black)]
273    #[case::green(0, 0b1000_0000, 0b11, Green)]
274    #[case::bright_green(0, 0b1100_0001, 0b11, BrightGreen)]
275    #[case::bright_yellow(0b1000_0000, 0b1100_0001, 0b11, BrightYellow)]
276    #[case::bright_white(0xFF, 0xFF, 0xFF, BrightWhite)]
277    fn test_color_choice_ansi_from_rgb(
278        #[case] r: u8,
279        #[case] g: u8,
280        #[case] b: u8,
281        #[case] expected: AnsiColors,
282    ) {
283        assert_eq!(expected, ColorChoice::ansi_from_rgb(r, g, b));
284    }
285}
fancy_tree/color/choice.rs

fancy_tree/color/
choice.rs
