mycelium_trace/
color.rs

1use core::fmt;
2
3use crate::writer::MakeWriter;
4pub trait SetColor {
5    fn set_fg_color(&mut self, color: Color);
6    fn fg_color(&self) -> Color;
7
8    /// Sets bold text.
9    ///
10    /// This may brighten a text color if bold text is not supported.
11    fn set_bold(&mut self, bold: bool);
12
13    fn with_bold(&mut self) -> WithBold<'_, Self>
14    where
15        Self: fmt::Write + Sized,
16    {
17        self.set_bold(true);
18        WithBold { writer: self }
19    }
20
21    #[must_use]
22    fn with_fg_color(&mut self, color: Color) -> WithFgColor<'_, Self>
23    where
24        Self: fmt::Write + Sized,
25    {
26        let prev_color = self.fg_color();
27        self.set_fg_color(color);
28        WithFgColor {
29            writer: self,
30            prev_color,
31        }
32    }
33}
34
35#[derive(Copy, Clone, Debug, Eq, PartialEq)]
36#[repr(usize)]
37pub enum Color {
38    Black = 0,
39    Red,
40    Green,
41    Yellow,
42    Blue,
43    Magenta,
44    Cyan,
45    White,
46    Default,
47    BrightBlack,
48    BrightRed,
49    BrightGreen,
50    BrightYellow,
51    BrightBlue,
52    BrightMagenta,
53    BrightCyan,
54    BrightWhite,
55}
56
57#[derive(Debug, Eq, PartialEq)]
58pub struct WithFgColor<'writer, W>
59where
60    W: fmt::Write + SetColor,
61{
62    writer: &'writer mut W,
63    prev_color: Color,
64}
65
66#[derive(Debug, Eq, PartialEq)]
67pub struct WithBold<'writer, W>
68where
69    W: fmt::Write + SetColor,
70{
71    writer: &'writer mut W,
72}
73
74#[derive(Copy, Clone, Debug, Eq, PartialEq)]
75pub struct AnsiEscapes<W> {
76    writer: W,
77    current_fg: Color,
78}
79
80impl<W: SetColor> SetColor for &'_ mut W {
81    #[inline]
82    fn set_fg_color(&mut self, color: Color) {
83        W::set_fg_color(self, color)
84    }
85
86    #[inline]
87    fn fg_color(&self) -> Color {
88        W::fg_color(self)
89    }
90
91    #[inline]
92    fn set_bold(&mut self, bold: bool) {
93        W::set_bold(self, bold)
94    }
95}
96
97// === impl WithFgColor ===
98
99impl<W> fmt::Write for WithFgColor<'_, W>
100where
101    W: fmt::Write + SetColor,
102{
103    #[inline]
104    fn write_str(&mut self, s: &str) -> fmt::Result {
105        self.writer.write_str(s)
106    }
107
108    #[inline]
109    fn write_char(&mut self, c: char) -> fmt::Result {
110        self.writer.write_char(c)
111    }
112
113    #[inline]
114    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
115        self.writer.write_fmt(args)
116    }
117}
118
119impl<W> Drop for WithFgColor<'_, W>
120where
121    W: fmt::Write + SetColor,
122{
123    fn drop(&mut self) {
124        self.writer.set_fg_color(self.prev_color);
125    }
126}
127
128// === impl WithBold ===
129
130impl<W> fmt::Write for WithBold<'_, W>
131where
132    W: fmt::Write + SetColor,
133{
134    #[inline]
135    fn write_str(&mut self, s: &str) -> fmt::Result {
136        self.writer.write_str(s)
137    }
138
139    #[inline]
140    fn write_char(&mut self, c: char) -> fmt::Result {
141        self.writer.write_char(c)
142    }
143
144    #[inline]
145    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
146        self.writer.write_fmt(args)
147    }
148}
149
150impl<W> Drop for WithBold<'_, W>
151where
152    W: fmt::Write + SetColor,
153{
154    fn drop(&mut self) {
155        self.writer.set_bold(false);
156    }
157}
158
159// === impl AnsiEscapes ===
160
161impl<W> AnsiEscapes<W> {
162    pub fn new(writer: W) -> Self {
163        Self {
164            writer,
165            current_fg: Color::Default,
166        }
167    }
168}
169
170impl<'mk, W> MakeWriter<'mk> for AnsiEscapes<W>
171where
172    W: MakeWriter<'mk>,
173{
174    type Writer = AnsiEscapes<W::Writer>;
175
176    fn make_writer(&'mk self) -> Self::Writer {
177        AnsiEscapes {
178            writer: self.writer.make_writer(),
179            current_fg: self.current_fg,
180        }
181    }
182
183    fn enabled(&self, meta: &tracing_core::Metadata<'_>) -> bool {
184        self.writer.enabled(meta)
185    }
186
187    /// Returns a [`Writer`] for writing data from the span or event described
188    /// by the provided [`Metadata`].
189    ///
190    /// By default, this calls [`self.make_writer()`][make_writer], ignoring
191    /// the provided metadata, but implementations can override this to provide
192    /// metadata-specific behaviors.
193    ///
194    /// This method allows `MakeWriter` implementations to implement different
195    /// behaviors based on the span or event being written. The `MakeWriter`
196    /// type might return different writers based on the provided metadata, or
197    /// might write some values to the writer before or after providing it to
198    /// the caller.
199    ///
200    /// [`Writer`]: MakeWriter::Writer
201    /// [`Metadata`]: tracing_core::Metadata
202    /// [make_writer]: MakeWriter::make_writer
203    /// [`WARN`]: tracing_core::Level::WARN
204    /// [`ERROR`]: tracing_core::Level::ERROR
205    #[inline]
206    fn make_writer_for(&'mk self, meta: &tracing_core::Metadata<'_>) -> Option<Self::Writer> {
207        self.writer.make_writer_for(meta).map(|writer| AnsiEscapes {
208            writer,
209            current_fg: self.current_fg,
210        })
211    }
212
213    fn line_len(&self) -> usize {
214        self.writer.line_len()
215    }
216}
217
218impl<W> fmt::Write for AnsiEscapes<W>
219where
220    W: fmt::Write,
221{
222    #[inline]
223    fn write_str(&mut self, s: &str) -> fmt::Result {
224        self.writer.write_str(s)
225    }
226
227    #[inline]
228    fn write_char(&mut self, c: char) -> fmt::Result {
229        self.writer.write_char(c)
230    }
231
232    #[inline]
233    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
234        self.writer.write_fmt(args)
235    }
236}
237
238impl<W> AnsiEscapes<W> {
239    const ANSI_FG_COLOR_TABLE: [&'static str; 17] = [
240        "30", // black
241        "31", // red
242        "32", // green
243        "33", // yellow
244        "34", // blue
245        "35", // magenta
246        "36", // cyan
247        "37", // white
248        "39", // default
249        "90", // bright black
250        "91", // bright red
251        "92", // bright green
252        "93", // bright yellow
253        "94", // bright blue
254        "95", // bright magenta
255        "96", // bright cyan
256        "97", // bright white
257    ];
258
259    fn fg_code(&self) -> &'static str {
260        Self::ANSI_FG_COLOR_TABLE[self.current_fg as usize]
261    }
262}
263
264impl<W: fmt::Write> SetColor for AnsiEscapes<W> {
265    fn set_fg_color(&mut self, color: Color) {
266        self.current_fg = color;
267        let _ = write!(self.writer, "\x1b[{}m", self.fg_code());
268    }
269
270    fn fg_color(&self) -> Color {
271        self.current_fg
272    }
273
274    fn set_bold(&mut self, bold: bool) {
275        let _ = if bold {
276            self.writer.write_str("\x1b[1m")
277        } else {
278            self.writer.write_str("\x1b[22m")
279        };
280    }
281}