spdlog/sink/
std_stream_sink.rs

1//! Provides a std stream sink.
2
3use std::{
4    convert::Infallible,
5    io::{self, Write},
6    // Import `str` module for function `std::str::from_utf8`, because method `str::from_utf8` is
7    // stabilized since Rust 1.87.
8    //
9    // TODO: Remove this import when our MSRV reaches Rust 1.87.
10    str,
11};
12
13use crate::{
14    formatter::{Formatter, FormatterContext},
15    sink::{GetSinkProp, Sink, SinkProp},
16    sync::*,
17    terminal_style::{LevelStyles, Style, StyleMode},
18    Error, ErrorHandler, Level, LevelFilter, Record, Result, StringBuf,
19};
20
21/// An enum representing the available standard streams.
22#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
23pub enum StdStream {
24    /// Standard output.
25    Stdout,
26    /// Standard error.
27    Stderr,
28}
29
30impl StdStream {
31    fn via_write(&self) -> StdStreamDest<io::StdoutLock<'_>, io::StderrLock<'_>> {
32        match self {
33            Self::Stdout => StdStreamDest::Stdout(io::stdout().lock()),
34            Self::Stderr => StdStreamDest::Stderr(io::stderr().lock()),
35        }
36    }
37
38    fn via_macro(&self) -> StdStreamDest<via_macro::Stdout, via_macro::Stderr> {
39        match self {
40            Self::Stdout => StdStreamDest::Stdout(via_macro::Stdout),
41            Self::Stderr => StdStreamDest::Stderr(via_macro::Stderr),
42        }
43    }
44}
45
46// `io::stdout()` and `io::stderr()` return different types, and
47// `Std***::lock()` is not in any trait, so we need this struct to abstract
48// them.
49#[derive(Debug)]
50enum StdStreamDest<O, E> {
51    Stdout(O),
52    Stderr(E),
53}
54
55impl<O, E> StdStreamDest<O, E> {
56    #[allow(dead_code)]
57    fn stream_type(&self) -> StdStream {
58        match self {
59            Self::Stdout(_) => StdStream::Stdout,
60            Self::Stderr(_) => StdStream::Stderr,
61        }
62    }
63}
64
65macro_rules! impl_write_for_dest {
66    ( $dest:ty ) => {
67        impl Write for $dest {
68            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
69                match self {
70                    StdStreamDest::Stdout(stream) => stream.write(buf),
71                    StdStreamDest::Stderr(stream) => stream.write(buf),
72                }
73            }
74
75            fn flush(&mut self) -> io::Result<()> {
76                match self {
77                    StdStreamDest::Stdout(stream) => stream.flush(),
78                    StdStreamDest::Stderr(stream) => stream.flush(),
79                }
80            }
81        }
82    };
83}
84impl_write_for_dest!(StdStreamDest<io::StdoutLock<'_>, io::StderrLock<'_>>);
85impl_write_for_dest!(StdStreamDest<via_macro::Stdout, via_macro::Stderr>);
86
87mod via_macro {
88    use super::*;
89
90    fn bytes_to_str(buf: &[u8]) -> io::Result<&str> {
91        str::from_utf8(buf).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
92    }
93
94    pub(crate) struct Stdout;
95
96    impl Write for Stdout {
97        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
98            print!("{}", bytes_to_str(buf)?);
99            Ok(buf.len())
100        }
101
102        fn flush(&mut self) -> io::Result<()> {
103            io::stdout().flush()
104        }
105    }
106
107    pub(crate) struct Stderr;
108
109    impl Write for Stderr {
110        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
111            eprint!("{}", bytes_to_str(buf)?);
112            Ok(buf.len())
113        }
114
115        fn flush(&mut self) -> io::Result<()> {
116            io::stderr().flush()
117        }
118    }
119}
120
121/// A sink with a std stream as the target.
122///
123/// It writes styled text or plain text according to the given [`StyleMode`] and
124/// the current terminal environment.
125///
126/// Note that this sink always flushes the buffer once with each logging.
127pub struct StdStreamSink {
128    prop: SinkProp,
129    via_print_macro: bool,
130    std_stream: StdStream,
131    should_render_style: bool,
132    level_styles: LevelStyles,
133}
134
135impl StdStreamSink {
136    /// Gets a builder of `StdStreamSink` with default parameters:
137    ///
138    /// | Parameter         | Default Value                                                       |
139    /// |-------------------|---------------------------------------------------------------------|
140    /// | [level_filter]    | [`LevelFilter::All`]                                                |
141    /// | [formatter]       | [`FullFormatter`]                                                   |
142    /// | [error_handler]   | [`ErrorHandler::default()`]                                         |
143    /// |                   |                                                                     |
144    /// | [std_stream]      | *must be specified*                                                 |
145    /// | [style_mode]      | [`StyleMode::Auto`]                                                 |
146    /// | [via_print_macro] | `false`, or `true` if feature gate `std-stream-captured` is enabled |
147    ///
148    /// [level_filter]: StdStreamSinkBuilder::level_filter
149    /// [formatter]: StdStreamSinkBuilder::formatter
150    /// [`FullFormatter`]: crate::formatter::FullFormatter
151    /// [error_handler]: StdStreamSinkBuilder::error_handler
152    /// [std_stream]: StdStreamSinkBuilder::std_stream
153    /// [style_mode]: StdStreamSinkBuilder::style_mode
154    /// [via_print_macro]: StdStreamSinkBuilder::via_print_macro
155    #[must_use]
156    pub fn builder() -> StdStreamSinkBuilder<()> {
157        StdStreamSinkBuilder {
158            prop: SinkProp::default(),
159            std_stream: (),
160            style_mode: StyleMode::Auto,
161            via_print_macro: cfg!(feature = "std-stream-captured"),
162        }
163    }
164
165    /// Constructs a `StdStreamSink`.
166    #[deprecated(
167        since = "0.3.0",
168        note = "it may be removed in the future, use `StdStreamSink::builder()` instead"
169    )]
170    #[must_use]
171    pub fn new(std_stream: StdStream, style_mode: StyleMode) -> StdStreamSink {
172        Self::builder()
173            .std_stream(std_stream)
174            .style_mode(style_mode)
175            .build()
176            .unwrap()
177    }
178
179    /// Sets the style of the specified log level.
180    pub fn set_style(&mut self, level: Level, style: Style) {
181        self.level_styles.set_style(level, style);
182    }
183
184    /// Sets the style mode.
185    pub fn set_style_mode(&mut self, style_mode: StyleMode) {
186        self.should_render_style = Self::should_render_style(style_mode, self.std_stream);
187    }
188
189    #[must_use]
190    fn should_render_style(style_mode: StyleMode, stream: StdStream) -> bool {
191        use is_terminal::IsTerminal;
192        let is_terminal = match stream {
193            StdStream::Stdout => io::stdout().is_terminal(),
194            StdStream::Stderr => io::stderr().is_terminal(),
195        };
196
197        match style_mode {
198            StyleMode::Always => true,
199            StyleMode::Auto => is_terminal && enable_ansi_escape_sequences(),
200            StyleMode::Never => false,
201        }
202    }
203}
204
205impl GetSinkProp for StdStreamSink {
206    fn prop(&self) -> &SinkProp {
207        &self.prop
208    }
209}
210
211impl Sink for StdStreamSink {
212    fn log(&self, record: &Record) -> Result<()> {
213        let mut string_buf = StringBuf::new();
214        let mut ctx = FormatterContext::new();
215        self.prop
216            .formatter()
217            .format(record, &mut string_buf, &mut ctx)?;
218
219        if !self.via_print_macro {
220            self.log_write(record, &string_buf, &ctx, self.std_stream.via_write())
221        } else {
222            self.log_write(record, &string_buf, &ctx, self.std_stream.via_macro())
223        }
224    }
225
226    fn flush(&self) -> Result<()> {
227        if !self.via_print_macro {
228            self.std_stream.via_write().flush()
229        } else {
230            self.std_stream.via_macro().flush()
231        }
232        .map_err(Error::FlushBuffer)
233    }
234}
235
236impl StdStreamSink {
237    fn log_write<O: Write, E: Write>(
238        &self,
239        record: &Record,
240        string_buf: &StringBuf,
241        ctx: &FormatterContext<'_>,
242        mut dest: StdStreamDest<O, E>,
243    ) -> Result<()>
244    where
245        StdStreamDest<O, E>: Write,
246    {
247        (|| {
248            // TODO: Simplify the if block when our MSRV reaches let-chain support.
249            if self.should_render_style {
250                if let Some(style_range) = ctx.style_range() {
251                    let style = self.level_styles.style(record.level());
252
253                    dest.write_all(string_buf[..style_range.start].as_bytes())?;
254                    style.write_start(&mut dest)?;
255                    dest.write_all(string_buf[style_range.start..style_range.end].as_bytes())?;
256                    style.write_end(&mut dest)?;
257                    dest.write_all(string_buf[style_range.end..].as_bytes())?;
258                } else {
259                    dest.write_all(string_buf.as_bytes())?;
260                }
261            } else {
262                dest.write_all(string_buf.as_bytes())?;
263            }
264
265            Ok(())
266        })()
267        .map_err(Error::WriteRecord)?;
268
269        // stderr is not buffered, so we don't need to flush it.
270        // https://doc.rust-lang.org/std/io/fn.stderr.html
271        if let StdStream::Stdout = self.std_stream {
272            dest.flush().map_err(Error::FlushBuffer)?;
273        }
274
275        Ok(())
276    }
277}
278
279// --------------------------------------------------
280
281/// #
282#[doc = include_str!("../include/doc/generic-builder-note.md")]
283pub struct StdStreamSinkBuilder<ArgSS> {
284    prop: SinkProp,
285    std_stream: ArgSS,
286    style_mode: StyleMode,
287    via_print_macro: bool,
288}
289
290impl<ArgSS> StdStreamSinkBuilder<ArgSS> {
291    /// Specifies the target standard stream as stdout.
292    ///
293    /// This is equivalent to `std_stream(StdStream::Stdout)`.
294    #[must_use]
295    pub fn stdout(self) -> StdStreamSinkBuilder<StdStream> {
296        self.std_stream(StdStream::Stdout)
297    }
298
299    /// Specifies the target standard stream as stderr.
300    ///
301    /// This is equivalent to `std_stream(StdStream::Stderr)`.
302    #[must_use]
303    pub fn stderr(self) -> StdStreamSinkBuilder<StdStream> {
304        self.std_stream(StdStream::Stderr)
305    }
306
307    /// Specifies the target standard stream.
308    ///
309    /// This parameter is **required**.
310    #[must_use]
311    pub fn std_stream(self, std_stream: StdStream) -> StdStreamSinkBuilder<StdStream> {
312        StdStreamSinkBuilder {
313            prop: self.prop,
314            std_stream,
315            style_mode: self.style_mode,
316            via_print_macro: self.via_print_macro,
317        }
318    }
319
320    /// Specifies the style mode.
321    ///
322    /// This parameter is **optional**, and defaults to [StyleMode::Auto].
323    #[must_use]
324    pub fn style_mode(mut self, style_mode: StyleMode) -> Self {
325        self.style_mode = style_mode;
326        self
327    }
328
329    /// Specifies to use `print!` and `eprint!` macros for output.
330    ///
331    /// If enabled, the sink will use [`print!`] and [`eprint!`] macros instead
332    /// of [`io::Write`] trait with [`io::stdout`] and [`io::stderr`] to output
333    /// logs. This is useful if you want the logs to be [captured] by `cargo
334    /// test` and `cargo bench`.
335    ///
336    /// This parameter is **optional**, and defaults to `false`, or defaults to
337    /// `true` if feature gate `std-stream-captured` is enabled.
338    ///
339    /// A convienient way to enable it for `cargo test` and `cargo bench` is to
340    /// add the following lines to your `Cargo.toml`:
341    ///
342    /// ```toml
343    /// # Note that it's not [dependencies]
344    ///
345    /// [dev-dependencies]
346    /// spdlog-rs = { version = "...", features = ["std-stream-captured"] }
347    /// ```
348    ///
349    /// [captured]: https://doc.rust-lang.org/book/ch11-02-running-tests.html#showing-function-output
350    #[must_use]
351    pub fn via_print_macro(mut self) -> Self {
352        self.via_print_macro = true;
353        self
354    }
355
356    // Prop
357    //
358
359    /// Specifies a log level filter.
360    ///
361    /// This parameter is **optional**, and defaults to [`LevelFilter::All`].
362    #[must_use]
363    pub fn level_filter(self, level_filter: LevelFilter) -> Self {
364        self.prop.set_level_filter(level_filter);
365        self
366    }
367
368    /// Specifies a formatter.
369    ///
370    /// This parameter is **optional**, and defaults to [`FullFormatter`].
371    ///
372    /// [`FullFormatter`]: crate::formatter::FullFormatter
373    #[must_use]
374    pub fn formatter<F>(self, formatter: F) -> Self
375    where
376        F: Formatter + 'static,
377    {
378        self.prop.set_formatter(formatter);
379        self
380    }
381
382    /// Specifies an error handler.
383    ///
384    /// This parameter is **optional**, and defaults to
385    /// [`ErrorHandler::default()`].
386    #[must_use]
387    pub fn error_handler<F: Into<ErrorHandler>>(self, handler: F) -> Self {
388        self.prop.set_error_handler(handler);
389        self
390    }
391}
392
393impl StdStreamSinkBuilder<()> {
394    #[doc(hidden)]
395    #[deprecated(note = "\n\n\
396        builder compile-time error:\n\
397        - missing required parameter `std_stream`\n\n\
398    ")]
399    pub fn build(self, _: Infallible) {}
400
401    #[doc(hidden)]
402    #[deprecated(note = "\n\n\
403        builder compile-time error:\n\
404        - missing required parameter `std_stream`\n\n\
405    ")]
406    pub fn build_arc(self, _: Infallible) {}
407}
408
409impl StdStreamSinkBuilder<StdStream> {
410    /// Builds a [`StdStreamSink`].
411    pub fn build(self) -> Result<StdStreamSink> {
412        Ok(StdStreamSink {
413            prop: self.prop,
414            via_print_macro: self.via_print_macro,
415            std_stream: self.std_stream,
416            should_render_style: StdStreamSink::should_render_style(
417                self.style_mode,
418                self.std_stream,
419            ),
420            level_styles: LevelStyles::default(),
421        })
422    }
423
424    /// Builds a `Arc<StdStreamSink>`.
425    ///
426    /// This is a shorthand method for `.build().map(Arc::new)`.
427    pub fn build_arc(self) -> Result<Arc<StdStreamSink>> {
428        self.build().map(Arc::new)
429    }
430}
431
432// --------------------------------------------------
433#[cfg(windows)]
434#[must_use]
435fn enable_ansi_escape_sequences() -> bool {
436    #[must_use]
437    fn enable() -> bool {
438        use winapi::um::{
439            consoleapi::{GetConsoleMode, SetConsoleMode},
440            handleapi::INVALID_HANDLE_VALUE,
441            processenv::GetStdHandle,
442            winbase::STD_OUTPUT_HANDLE,
443            wincon::ENABLE_VIRTUAL_TERMINAL_PROCESSING,
444        };
445
446        let stdout_handle = unsafe { GetStdHandle(STD_OUTPUT_HANDLE) };
447        if stdout_handle == INVALID_HANDLE_VALUE {
448            return false;
449        }
450
451        let mut original_mode = 0;
452        if unsafe { GetConsoleMode(stdout_handle, &mut original_mode) } == 0 {
453            return false;
454        }
455
456        let new_mode = original_mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING;
457        (unsafe { SetConsoleMode(stdout_handle, new_mode) }) != 0
458    }
459
460    use once_cell::sync::OnceCell;
461
462    static INIT: OnceCell<bool> = OnceCell::new();
463
464    *INIT.get_or_init(enable)
465}
466
467#[cfg(not(windows))]
468#[must_use]
469fn enable_ansi_escape_sequences() -> bool {
470    true
471}