1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

//! Helpers to prevent progress bars from clobbering your console output.
use std::io;
use std::marker::PhantomData;

use indicatif::MultiProgress;
use tracing_subscriber::fmt::MakeWriter;

pub trait WriterTarget: private::Sealed {}

/// Marker for where the [IndicatifWriter] should write to.
pub struct Stdout {}
/// Marker for where the [IndicatifWriter] should write to.
pub struct Stderr {}

impl WriterTarget for Stdout {}
impl WriterTarget for Stderr {}

// TODO(emersonford): find a cleaner way to integrate this layer with fmt::Layer.
/// A wrapper around [std::io::stdout()] or [std::io::stderr()] to ensure that output to either is
/// not clobbered by active progress bars. This should be passed into tracing fmt's layer so
/// tracing log entries are not clobbered.
pub struct IndicatifWriter<Target = Stderr> {
    progress_bars: MultiProgress,
    inner: PhantomData<Target>,
}

impl<T> IndicatifWriter<T>
where
    T: WriterTarget,
{
    pub(crate) fn new(mp: MultiProgress) -> Self {
        Self {
            progress_bars: mp,
            inner: PhantomData,
        }
    }
}

impl<T> Clone for IndicatifWriter<T> {
    fn clone(&self) -> Self {
        Self {
            progress_bars: self.progress_bars.clone(),
            inner: self.inner,
        }
    }
}

impl io::Write for IndicatifWriter<Stdout> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.progress_bars.suspend(|| io::stdout().write(buf))
    }

    fn flush(&mut self) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stdout().flush())
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
        self.progress_bars
            .suspend(|| io::stdout().write_vectored(bufs))
    }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stdout().write_all(buf))
    }

    fn write_fmt(&mut self, fmt: std::fmt::Arguments<'_>) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stdout().write_fmt(fmt))
    }
}

impl io::Write for IndicatifWriter<Stderr> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.progress_bars.suspend(|| io::stderr().write(buf))
    }

    fn flush(&mut self) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stderr().flush())
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
        self.progress_bars
            .suspend(|| io::stderr().write_vectored(bufs))
    }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stderr().write_all(buf))
    }

    fn write_fmt(&mut self, fmt: std::fmt::Arguments<'_>) -> io::Result<()> {
        self.progress_bars.suspend(|| io::stderr().write_fmt(fmt))
    }
}

impl<'a> MakeWriter<'a> for IndicatifWriter<Stdout> {
    type Writer = IndicatifWriter<Stdout>;

    fn make_writer(&'a self) -> Self::Writer {
        self.clone()
    }
}

impl<'a> MakeWriter<'a> for IndicatifWriter<Stderr> {
    type Writer = IndicatifWriter<Stderr>;

    fn make_writer(&'a self) -> Self::Writer {
        self.clone()
    }
}

mod private {
    pub trait Sealed {}

    impl Sealed for super::Stdout {}
    impl Sealed for super::Stderr {}
}

/// Returns the stderr writer (equivalent to
/// [get_stderr_writer](crate::IndicatifLayer::get_stderr_writer)) of the registered
/// [IndicatifLayer](crate::IndicatifLayer) for the current default tracing subscriber.
///
/// Returns `None` if there is either no default tracing subscriber or if there is not a
/// `IndicatifLayer` registered with that subscriber.
pub fn get_indicatif_stderr_writer() -> Option<IndicatifWriter<Stderr>> {
    tracing::dispatcher::get_default(|dispatch| {
        dispatch
            .downcast_ref::<crate::WithStderrWriter>()
            .and_then(|ctx| {
                let mut ret: Option<IndicatifWriter<Stderr>> = None;
                ctx.with_context(dispatch, |writer| {
                    ret = Some(writer);
                });

                ret
            })
    })
}

/// Returns the stdout writer (equivalent to
/// [get_stdout_writer](crate::IndicatifLayer::get_stdout_writer)) of the registered
/// [IndicatifLayer](crate::IndicatifLayer) for the current default tracing subscriber.
///
/// Returns `None` if there is either no default tracing subscriber or if there is not a
/// `IndicatifLayer` registered with that subscriber.
pub fn get_indicatif_stdout_writer() -> Option<IndicatifWriter<Stdout>> {
    tracing::dispatcher::get_default(|dispatch| {
        dispatch
            .downcast_ref::<crate::WithStdoutWriter>()
            .and_then(|ctx| {
                let mut ret: Option<IndicatifWriter<Stdout>> = None;
                ctx.with_context(dispatch, |writer| {
                    ret = Some(writer);
                });

                ret
            })
    })
}