tracing_indicatif/
writer.rs

1//! Helpers to prevent progress bars from clobbering your console output.
2use std::io;
3use std::marker::PhantomData;
4
5use indicatif::MultiProgress;
6use tracing_subscriber::fmt::MakeWriter;
7
8pub trait WriterTarget: private::Sealed {}
9
10/// Marker for where the [`IndicatifWriter`] should write to.
11pub struct Stdout {}
12/// Marker for where the [`IndicatifWriter`] should write to.
13pub struct Stderr {}
14
15impl WriterTarget for Stdout {}
16impl WriterTarget for Stderr {}
17
18// TODO(emersonford): find a cleaner way to integrate this layer with fmt::Layer.
19/// A wrapper around [`std::io::stdout()`] or [`std::io::stderr()`] to ensure that output to either is
20/// not clobbered by active progress bars. This should be passed into tracing fmt's layer so
21/// tracing log entries are not clobbered.
22pub struct IndicatifWriter<Target = Stderr> {
23    progress_bars: MultiProgress,
24    inner: PhantomData<Target>,
25}
26
27impl<T> IndicatifWriter<T>
28where
29    T: WriterTarget,
30{
31    pub(crate) fn new(mp: MultiProgress) -> Self {
32        Self {
33            progress_bars: mp,
34            inner: PhantomData,
35        }
36    }
37}
38
39impl<T> Clone for IndicatifWriter<T> {
40    fn clone(&self) -> Self {
41        Self {
42            progress_bars: self.progress_bars.clone(),
43            inner: self.inner,
44        }
45    }
46}
47
48impl io::Write for IndicatifWriter<Stdout> {
49    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
50        self.progress_bars.suspend(|| io::stdout().write(buf))
51    }
52
53    fn flush(&mut self) -> io::Result<()> {
54        self.progress_bars.suspend(|| io::stdout().flush())
55    }
56
57    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
58        self.progress_bars
59            .suspend(|| io::stdout().write_vectored(bufs))
60    }
61
62    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
63        self.progress_bars.suspend(|| io::stdout().write_all(buf))
64    }
65
66    fn write_fmt(&mut self, fmt: std::fmt::Arguments<'_>) -> io::Result<()> {
67        self.progress_bars.suspend(|| io::stdout().write_fmt(fmt))
68    }
69}
70
71impl io::Write for IndicatifWriter<Stderr> {
72    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
73        self.progress_bars.suspend(|| io::stderr().write(buf))
74    }
75
76    fn flush(&mut self) -> io::Result<()> {
77        self.progress_bars.suspend(|| io::stderr().flush())
78    }
79
80    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
81        self.progress_bars
82            .suspend(|| io::stderr().write_vectored(bufs))
83    }
84
85    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
86        self.progress_bars.suspend(|| io::stderr().write_all(buf))
87    }
88
89    fn write_fmt(&mut self, fmt: std::fmt::Arguments<'_>) -> io::Result<()> {
90        self.progress_bars.suspend(|| io::stderr().write_fmt(fmt))
91    }
92}
93
94impl<'a> MakeWriter<'a> for IndicatifWriter<Stdout> {
95    type Writer = IndicatifWriter<Stdout>;
96
97    fn make_writer(&'a self) -> Self::Writer {
98        self.clone()
99    }
100}
101
102impl<'a> MakeWriter<'a> for IndicatifWriter<Stderr> {
103    type Writer = IndicatifWriter<Stderr>;
104
105    fn make_writer(&'a self) -> Self::Writer {
106        self.clone()
107    }
108}
109
110mod private {
111    pub trait Sealed {}
112
113    impl Sealed for super::Stdout {}
114    impl Sealed for super::Stderr {}
115}
116
117/// Returns the stderr writer (equivalent to
118/// [`get_stderr_writer`](crate::IndicatifLayer::get_stderr_writer)) of the registered
119/// [`IndicatifLayer`](crate::IndicatifLayer) for the current default tracing subscriber.
120///
121/// Returns `None` if there is either no default tracing subscriber or if there is not a
122/// `IndicatifLayer` registered with that subscriber.
123pub fn get_indicatif_stderr_writer() -> Option<IndicatifWriter<Stderr>> {
124    tracing::dispatcher::get_default(|dispatch| {
125        dispatch
126            .downcast_ref::<crate::WithStderrWriter>()
127            .and_then(|ctx| {
128                let mut ret: Option<IndicatifWriter<Stderr>> = None;
129                ctx.with_context(dispatch, |writer| {
130                    ret = Some(writer);
131                });
132
133                ret
134            })
135    })
136}
137
138/// Returns the stdout writer (equivalent to
139/// [`get_stdout_writer`](crate::IndicatifLayer::get_stdout_writer)) of the registered
140/// [`IndicatifLayer`](crate::IndicatifLayer) for the current default tracing subscriber.
141///
142/// Returns `None` if there is either no default tracing subscriber or if there is not a
143/// `IndicatifLayer` registered with that subscriber.
144pub fn get_indicatif_stdout_writer() -> Option<IndicatifWriter<Stdout>> {
145    tracing::dispatcher::get_default(|dispatch| {
146        dispatch
147            .downcast_ref::<crate::WithStdoutWriter>()
148            .and_then(|ctx| {
149                let mut ret: Option<IndicatifWriter<Stdout>> = None;
150                ctx.with_context(dispatch, |writer| {
151                    ret = Some(writer);
152                });
153
154                ret
155            })
156    })
157}