spdlog/sink/
std_stream_sink.rs

1//! Provides a std stream sink.
2
3use std::{
4    convert::Infallible,
5    io::{self, Write},
6};
7
8use crate::{
9    formatter::{Formatter, FormatterContext},
10    sink::{GetSinkProp, Sink, SinkProp},
11    terminal_style::{LevelStyles, Style, StyleMode},
12    Error, ErrorHandler, Level, LevelFilter, Record, Result, StringBuf,
13};
14
15/// An enum representing the available standard streams.
16#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
17pub enum StdStream {
18    /// Standard output.
19    Stdout,
20    /// Standard error.
21    Stderr,
22}
23
24// `io::stdout()` and `io::stderr()` return different types, and
25// `Std***::lock()` is not in any trait, so we need this struct to abstract
26// them.
27#[derive(Debug)]
28enum StdStreamDest<O, E> {
29    Stdout(O),
30    Stderr(E),
31}
32
33impl StdStreamDest<io::Stdout, io::Stderr> {
34    #[must_use]
35    fn new(stream: StdStream) -> Self {
36        match stream {
37            StdStream::Stdout => StdStreamDest::Stdout(io::stdout()),
38            StdStream::Stderr => StdStreamDest::Stderr(io::stderr()),
39        }
40    }
41
42    #[must_use]
43    fn lock(&self) -> StdStreamDest<io::StdoutLock<'_>, io::StderrLock<'_>> {
44        match self {
45            StdStreamDest::Stdout(stream) => StdStreamDest::Stdout(stream.lock()),
46            StdStreamDest::Stderr(stream) => StdStreamDest::Stderr(stream.lock()),
47        }
48    }
49
50    fn stream_type(&self) -> StdStream {
51        match self {
52            StdStreamDest::Stdout(_) => StdStream::Stdout,
53            StdStreamDest::Stderr(_) => StdStream::Stderr,
54        }
55    }
56}
57
58macro_rules! impl_write_for_dest {
59    ( $dest:ty ) => {
60        impl Write for $dest {
61            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
62                match self {
63                    StdStreamDest::Stdout(stream) => stream.write(buf),
64                    StdStreamDest::Stderr(stream) => stream.write(buf),
65                }
66            }
67
68            fn flush(&mut self) -> io::Result<()> {
69                match self {
70                    StdStreamDest::Stdout(stream) => stream.flush(),
71                    StdStreamDest::Stderr(stream) => stream.flush(),
72                }
73            }
74        }
75    };
76}
77impl_write_for_dest!(StdStreamDest<io::Stdout, io::Stderr>);
78impl_write_for_dest!(StdStreamDest<io::StdoutLock<'_>, io::StderrLock<'_>>);
79
80/// A sink with a std stream as the target.
81///
82/// It writes styled text or plain text according to the given [`StyleMode`] and
83/// the current terminal environment.
84///
85/// Note that this sink always flushes the buffer once with each logging.
86pub struct StdStreamSink {
87    prop: SinkProp,
88    dest: StdStreamDest<io::Stdout, io::Stderr>,
89    should_render_style: bool,
90    level_styles: LevelStyles,
91}
92
93impl StdStreamSink {
94    /// Gets a builder of `StdStreamSink` with default parameters:
95    ///
96    /// | Parameter         | Default Value               |
97    /// |-------------------|-----------------------------|
98    /// | [level_filter]    | `All`                       |
99    /// | [formatter]       | `FullFormatter`             |
100    /// | [error_handler]   | [`ErrorHandler::default()`] |
101    /// |                   |                             |
102    /// | [std_stream]      | *must be specified*         |
103    /// | [style_mode]      | `Auto`                      |
104    ///
105    /// [level_filter]: StdStreamSinkBuilder::level_filter
106    /// [formatter]: StdStreamSinkBuilder::formatter
107    /// [error_handler]: StdStreamSinkBuilder::error_handler
108    /// [`ErrorHandler::default()`]: crate::error::ErrorHandler::default()
109    /// [std_stream]: StdStreamSinkBuilder::std_stream
110    /// [style_mode]: StdStreamSinkBuilder::style_mode
111    #[must_use]
112    pub fn builder() -> StdStreamSinkBuilder<()> {
113        StdStreamSinkBuilder {
114            prop: SinkProp::default(),
115            std_stream: (),
116            style_mode: StyleMode::Auto,
117        }
118    }
119
120    /// Constructs a `StdStreamSink`.
121    #[deprecated(
122        since = "0.3.0",
123        note = "it may be removed in the future, use `StdStreamSink::builder()` instead"
124    )]
125    #[must_use]
126    pub fn new(std_stream: StdStream, style_mode: StyleMode) -> StdStreamSink {
127        Self::builder()
128            .std_stream(std_stream)
129            .style_mode(style_mode)
130            .build()
131            .unwrap()
132    }
133
134    /// Sets the style of the specified log level.
135    pub fn set_style(&mut self, level: Level, style: Style) {
136        self.level_styles.set_style(level, style);
137    }
138
139    /// Sets the style mode.
140    pub fn set_style_mode(&mut self, style_mode: StyleMode) {
141        self.should_render_style = Self::should_render_style(style_mode, self.dest.stream_type());
142    }
143
144    #[must_use]
145    fn should_render_style(style_mode: StyleMode, stream: StdStream) -> bool {
146        use is_terminal::IsTerminal;
147        let is_terminal = match stream {
148            StdStream::Stdout => io::stdout().is_terminal(),
149            StdStream::Stderr => io::stderr().is_terminal(),
150        };
151
152        match style_mode {
153            StyleMode::Always => true,
154            StyleMode::Auto => is_terminal && enable_ansi_escape_sequences(),
155            StyleMode::Never => false,
156        }
157    }
158}
159
160impl GetSinkProp for StdStreamSink {
161    fn prop(&self) -> &SinkProp {
162        &self.prop
163    }
164}
165
166impl Sink for StdStreamSink {
167    fn log(&self, record: &Record) -> Result<()> {
168        let mut string_buf = StringBuf::new();
169        let mut ctx = FormatterContext::new();
170        self.prop
171            .formatter()
172            .format(record, &mut string_buf, &mut ctx)?;
173
174        let mut dest = self.dest.lock();
175
176        (|| {
177            // TODO: Simplify the if block when our MSRV reaches let-chain support.
178            if self.should_render_style {
179                if let Some(style_range) = ctx.style_range() {
180                    let style = self.level_styles.style(record.level());
181
182                    dest.write_all(string_buf[..style_range.start].as_bytes())?;
183                    style.write_start(&mut dest)?;
184                    dest.write_all(string_buf[style_range.start..style_range.end].as_bytes())?;
185                    style.write_end(&mut dest)?;
186                    dest.write_all(string_buf[style_range.end..].as_bytes())?;
187                } else {
188                    dest.write_all(string_buf.as_bytes())?;
189                }
190            } else {
191                dest.write_all(string_buf.as_bytes())?;
192            }
193
194            Ok(())
195        })()
196        .map_err(Error::WriteRecord)?;
197
198        // stderr is not buffered, so we don't need to flush it.
199        // https://doc.rust-lang.org/std/io/fn.stderr.html
200        if let StdStreamDest::Stdout(_) = dest {
201            dest.flush().map_err(Error::FlushBuffer)?;
202        }
203
204        Ok(())
205    }
206
207    fn flush(&self) -> Result<()> {
208        self.dest.lock().flush().map_err(Error::FlushBuffer)
209    }
210}
211
212// --------------------------------------------------
213
214/// #
215#[doc = include_str!("../include/doc/generic-builder-note.md")]
216pub struct StdStreamSinkBuilder<ArgSS> {
217    prop: SinkProp,
218    std_stream: ArgSS,
219    style_mode: StyleMode,
220}
221
222impl<ArgSS> StdStreamSinkBuilder<ArgSS> {
223    /// Specifies the target standard stream as stdout.
224    ///
225    /// This is equivalent to `std_stream(StdStream::Stdout)`.
226    #[must_use]
227    pub fn stdout(self) -> StdStreamSinkBuilder<StdStream> {
228        self.std_stream(StdStream::Stdout)
229    }
230
231    /// Specifies the target standard stream as stderr.
232    ///
233    /// This is equivalent to `std_stream(StdStream::Stderr)`.
234    #[must_use]
235    pub fn stderr(self) -> StdStreamSinkBuilder<StdStream> {
236        self.std_stream(StdStream::Stderr)
237    }
238
239    /// Specifies the target standard stream.
240    ///
241    /// This parameter is **required**.
242    #[must_use]
243    pub fn std_stream(self, std_stream: StdStream) -> StdStreamSinkBuilder<StdStream> {
244        StdStreamSinkBuilder {
245            prop: self.prop,
246            std_stream,
247            style_mode: self.style_mode,
248        }
249    }
250
251    /// Specifies the style mode.
252    ///
253    /// This parameter is **optional**.
254    #[must_use]
255    pub fn style_mode(mut self, style_mode: StyleMode) -> Self {
256        self.style_mode = style_mode;
257        self
258    }
259
260    // Prop
261    //
262
263    /// Specifies a log level filter.
264    ///
265    /// This parameter is **optional**.
266    #[must_use]
267    pub fn level_filter(self, level_filter: LevelFilter) -> Self {
268        self.prop.set_level_filter(level_filter);
269        self
270    }
271
272    /// Specifies a formatter.
273    ///
274    /// This parameter is **optional**.
275    #[must_use]
276    pub fn formatter<F>(self, formatter: F) -> Self
277    where
278        F: Formatter + 'static,
279    {
280        self.prop.set_formatter(formatter);
281        self
282    }
283
284    /// Specifies an error handler.
285    ///
286    /// This parameter is **optional**.
287    #[must_use]
288    pub fn error_handler<F: Into<ErrorHandler>>(self, handler: F) -> Self {
289        self.prop.set_error_handler(handler);
290        self
291    }
292}
293
294impl StdStreamSinkBuilder<()> {
295    #[doc(hidden)]
296    #[deprecated(note = "\n\n\
297        builder compile-time error:\n\
298        - missing required parameter `std_stream`\n\n\
299    ")]
300    pub fn build(self, _: Infallible) {}
301}
302
303impl StdStreamSinkBuilder<StdStream> {
304    /// Builds a [`StdStreamSink`].
305    pub fn build(self) -> Result<StdStreamSink> {
306        Ok(StdStreamSink {
307            prop: self.prop,
308            dest: StdStreamDest::new(self.std_stream),
309            should_render_style: StdStreamSink::should_render_style(
310                self.style_mode,
311                self.std_stream,
312            ),
313            level_styles: LevelStyles::default(),
314        })
315    }
316}
317
318// --------------------------------------------------
319#[cfg(windows)]
320#[must_use]
321fn enable_ansi_escape_sequences() -> bool {
322    #[must_use]
323    fn enable() -> bool {
324        use winapi::um::{
325            consoleapi::{GetConsoleMode, SetConsoleMode},
326            handleapi::INVALID_HANDLE_VALUE,
327            processenv::GetStdHandle,
328            winbase::STD_OUTPUT_HANDLE,
329            wincon::ENABLE_VIRTUAL_TERMINAL_PROCESSING,
330        };
331
332        let stdout_handle = unsafe { GetStdHandle(STD_OUTPUT_HANDLE) };
333        if stdout_handle == INVALID_HANDLE_VALUE {
334            return false;
335        }
336
337        let mut original_mode = 0;
338        if unsafe { GetConsoleMode(stdout_handle, &mut original_mode) } == 0 {
339            return false;
340        }
341
342        let new_mode = original_mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING;
343        (unsafe { SetConsoleMode(stdout_handle, new_mode) }) != 0
344    }
345
346    use once_cell::sync::OnceCell;
347
348    static INIT: OnceCell<bool> = OnceCell::new();
349
350    *INIT.get_or_init(enable)
351}
352
353#[cfg(not(windows))]
354#[must_use]
355fn enable_ansi_escape_sequences() -> bool {
356    true
357}