color_ansi/writer/
mod.rs

1use crate::{AnsiColor, AnsiStyle};
2use std::{fmt::Debug, io::Write};
3
4/// The ability of a writer to support ANSI escape codes.
5#[derive(Copy, Clone, Debug)]
6pub enum AnsiAbility {
7    /// Devices that do not support any color.
8    Disabled,
9    /// Old Windows devices,
10    ///
11    /// **Note**: this does not represent the windows platform, modern windows platforms also support [`AnsiAbility::Full`]
12    Windows,
13    /// Devices that support all ansi colors.
14    Full,
15    /// The writer may or may not support ANSI escape codes, but it is
16    Auto,
17}
18
19impl Default for AnsiAbility {
20    fn default() -> Self {
21        Self::Full
22    }
23}
24
25/// A writer that supports ANSI escape codes.
26pub struct AnsiWriter<W> {
27    ability: AnsiAbility,
28    writer: W,
29}
30
31impl<W> Debug for AnsiWriter<W> {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "AnsiWriter({:?})", self.ability)
34    }
35}
36
37impl<W: Write> Write for AnsiWriter<W> {
38    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
39        self.writer.write(buf)
40    }
41    fn flush(&mut self) -> std::io::Result<()> {
42        self.writer.flush()
43    }
44}
45
46impl<W: Write> AnsiWriter<W> {
47    /// Create a new `AnsiWriter` with the given writer.
48    pub fn new(writer: W) -> Self {
49        Self { ability: AnsiAbility::default(), writer }
50    }
51    /// Create a new `AnsiWriter` with the given writer and ability.
52    pub fn get_ability(&self) -> AnsiAbility {
53        self.ability
54    }
55    /// Set the ability of the writer.
56    pub fn set_ability(&mut self, ability: AnsiAbility) {
57        self.ability = ability;
58    }
59    /// Create a new `AnsiWriter` with the given writer and ability.
60    pub fn with_ability(mut self, ability: AnsiAbility) -> Self {
61        self.ability = ability;
62        self
63    }
64    /// Set the style of the writer.
65    pub fn set_style(&mut self, style: &AnsiStyle) -> std::io::Result<()> {
66        if style.reset {
67            self.reset_style()?;
68        }
69        if style.bold {
70            self.write_str("\x1B[1m")?;
71        }
72        if style.dimmed {
73            self.write_str("\x1B[2m")?;
74        }
75        if style.italic {
76            self.write_str("\x1B[3m")?;
77        }
78        if style.underline {
79            self.write_str("\x1B[4m")?;
80        }
81        if style.strikethrough {
82            self.write_str("\x1B[9m")?;
83        }
84        if let Some(c) = &style.fg_color {
85            self.write_color(true, c, style.intense)?;
86        }
87        if let Some(c) = &style.bg_color {
88            self.write_color(false, c, style.intense)?;
89        }
90        Ok(())
91    }
92    /// Reset the style of the writer.
93    #[inline]
94    pub fn reset_style(&mut self) -> std::io::Result<()> {
95        self.write_str("\x1B[0m")
96    }
97    fn write_str(&mut self, s: &str) -> std::io::Result<()> {
98        self.write_all(s.as_bytes())
99    }
100
101    fn write_color(&mut self, fg: bool, c: &AnsiColor, intense: bool) -> std::io::Result<()> {
102        macro_rules! write_intense {
103            ($clr:expr) => {
104                if fg {
105                    self.write_str(concat!("\x1B[38;5;", $clr, "m"))
106                }
107                else {
108                    self.write_str(concat!("\x1B[48;5;", $clr, "m"))
109                }
110            };
111        }
112        macro_rules! write_normal {
113            ($clr:expr) => {
114                if fg { self.write_str(concat!("\x1B[3", $clr, "m")) } else { self.write_str(concat!("\x1B[4", $clr, "m")) }
115            };
116        }
117        macro_rules! write_var_ansi_code {
118            ($pre:expr, $($code:expr),+) => {{
119                // The loop generates at worst a literal of the form
120                // '255,255,255m' which is 12-bytes.
121                // The largest `pre` expression we currently use is 7 bytes.
122                // This gives us the maximum of 19-bytes for our work buffer.
123                let pre_len = $pre.len();
124                assert!(pre_len <= 7);
125                let mut fmt = [0u8; 19];
126                fmt[..pre_len].copy_from_slice($pre);
127                let mut i = pre_len - 1;
128                $(
129                    let c1: u8 = ($code / 100) % 10;
130                    let c2: u8 = ($code / 10) % 10;
131                    let c3: u8 = $code % 10;
132                    let mut printed = false;
133
134                    if c1 != 0 {
135                        printed = true;
136                        i += 1;
137                        fmt[i] = b'0' + c1;
138                    }
139                    if c2 != 0 || printed {
140                        i += 1;
141                        fmt[i] = b'0' + c2;
142                    }
143                    // If we received a zero value we must still print a value.
144                    i += 1;
145                    fmt[i] = b'0' + c3;
146                    i += 1;
147                    fmt[i] = b';';
148                )+
149
150                fmt[i] = b'm';
151                self.write_all(&fmt[0..i+1])
152            }}
153        }
154        macro_rules! write_custom {
155            ($ansi256:expr) => {
156                if fg { write_var_ansi_code!(b"\x1B[38;5;", $ansi256) } else { write_var_ansi_code!(b"\x1B[48;5;", $ansi256) }
157            };
158
159            ($r:expr, $g:expr, $b:expr) => {{
160                if fg {
161                    write_var_ansi_code!(b"\x1B[38;2;", $r, $g, $b)
162                }
163                else {
164                    write_var_ansi_code!(b"\x1B[48;2;", $r, $g, $b)
165                }
166            }};
167        }
168        if intense {
169            match *c {
170                AnsiColor::Black => write_intense!("8"),
171                AnsiColor::Blue => write_intense!("12"),
172                AnsiColor::Green => write_intense!("10"),
173                AnsiColor::Red => write_intense!("9"),
174                AnsiColor::Cyan => write_intense!("14"),
175                AnsiColor::Magenta => write_intense!("13"),
176                AnsiColor::Yellow => write_intense!("11"),
177                AnsiColor::White => write_intense!("15"),
178                AnsiColor::Ansi256(c) => write_custom!(c),
179                AnsiColor::Rgb(r, g, b) => write_custom!(r, g, b),
180            }
181        }
182        else {
183            match *c {
184                AnsiColor::Black => write_normal!("0"),
185                AnsiColor::Blue => write_normal!("4"),
186                AnsiColor::Green => write_normal!("2"),
187                AnsiColor::Red => write_normal!("1"),
188                AnsiColor::Cyan => write_normal!("6"),
189                AnsiColor::Magenta => write_normal!("5"),
190                AnsiColor::Yellow => write_normal!("3"),
191                AnsiColor::White => write_normal!("7"),
192                AnsiColor::Ansi256(c) => write_custom!(c),
193                AnsiColor::Rgb(r, g, b) => write_custom!(r, g, b),
194            }
195        }
196    }
197}