anstream/
auto.rs

1use crate::stream::AsLockedWrite;
2use crate::stream::RawStream;
3use crate::ColorChoice;
4use crate::StripStream;
5#[cfg(all(windows, feature = "wincon"))]
6use crate::WinconStream;
7
8/// [`std::io::Write`] that adapts ANSI escape codes to the underlying `Write`s capabilities
9#[derive(Debug)]
10pub struct AutoStream<S: RawStream> {
11    inner: StreamInner<S>,
12}
13
14#[derive(Debug)]
15enum StreamInner<S: RawStream> {
16    PassThrough(S),
17    Strip(StripStream<S>),
18    #[cfg(all(windows, feature = "wincon"))]
19    Wincon(WinconStream<S>),
20}
21
22impl<S> AutoStream<S>
23where
24    S: RawStream,
25{
26    /// Runtime control over styling behavior
27    #[inline]
28    pub fn new(raw: S, choice: ColorChoice) -> Self {
29        match choice {
30            #[cfg(feature = "auto")]
31            ColorChoice::Auto => Self::auto(raw),
32            #[cfg(not(feature = "auto"))]
33            ColorChoice::Auto => Self::never(raw),
34            ColorChoice::AlwaysAnsi => Self::always_ansi(raw),
35            ColorChoice::Always => Self::always(raw),
36            ColorChoice::Never => Self::never(raw),
37        }
38    }
39
40    /// Auto-adapt for the stream's capabilities
41    #[cfg(feature = "auto")]
42    #[inline]
43    pub fn auto(raw: S) -> Self {
44        let choice = Self::choice(&raw);
45        debug_assert_ne!(choice, ColorChoice::Auto);
46        Self::new(raw, choice)
47    }
48
49    /// Report the desired choice for the given stream
50    #[cfg(feature = "auto")]
51    pub fn choice(raw: &S) -> ColorChoice {
52        choice(raw)
53    }
54
55    /// Force ANSI escape codes to be passed through as-is, no matter what the inner `Write`
56    /// supports.
57    #[inline]
58    pub fn always_ansi(raw: S) -> Self {
59        #[cfg(feature = "auto")]
60        {
61            if raw.is_terminal() {
62                let _ = anstyle_query::windows::enable_ansi_colors();
63            }
64        }
65        Self::always_ansi_(raw)
66    }
67
68    #[inline]
69    fn always_ansi_(raw: S) -> Self {
70        let inner = StreamInner::PassThrough(raw);
71        AutoStream { inner }
72    }
73
74    /// Force color, no matter what the inner `Write` supports.
75    #[inline]
76    pub fn always(raw: S) -> Self {
77        if cfg!(windows) {
78            #[cfg(feature = "auto")]
79            let use_wincon = raw.is_terminal()
80                && !anstyle_query::windows::enable_ansi_colors().unwrap_or(true)
81                && !anstyle_query::term_supports_ansi_color();
82            #[cfg(not(feature = "auto"))]
83            let use_wincon = true;
84            if use_wincon {
85                Self::wincon(raw).unwrap_or_else(|raw| Self::always_ansi_(raw))
86            } else {
87                Self::always_ansi_(raw)
88            }
89        } else {
90            Self::always_ansi(raw)
91        }
92    }
93
94    /// Only pass printable data to the inner `Write`.
95    #[inline]
96    pub fn never(raw: S) -> Self {
97        let inner = StreamInner::Strip(StripStream::new(raw));
98        AutoStream { inner }
99    }
100
101    #[inline]
102    fn wincon(raw: S) -> Result<Self, S> {
103        #[cfg(all(windows, feature = "wincon"))]
104        {
105            Ok(Self {
106                inner: StreamInner::Wincon(WinconStream::new(raw)),
107            })
108        }
109        #[cfg(not(all(windows, feature = "wincon")))]
110        {
111            Err(raw)
112        }
113    }
114
115    /// Get the wrapped [`RawStream`]
116    #[inline]
117    pub fn into_inner(self) -> S {
118        match self.inner {
119            StreamInner::PassThrough(w) => w,
120            StreamInner::Strip(w) => w.into_inner(),
121            #[cfg(all(windows, feature = "wincon"))]
122            StreamInner::Wincon(w) => w.into_inner(),
123        }
124    }
125
126    #[inline]
127    pub fn is_terminal(&self) -> bool {
128        match &self.inner {
129            StreamInner::PassThrough(w) => w.is_terminal(),
130            StreamInner::Strip(w) => w.is_terminal(),
131            #[cfg(all(windows, feature = "wincon"))]
132            StreamInner::Wincon(_) => true, // its only ever a terminal
133        }
134    }
135
136    /// Prefer [`AutoStream::choice`]
137    ///
138    /// This doesn't report what is requested but what is currently active.
139    #[inline]
140    #[cfg(feature = "auto")]
141    pub fn current_choice(&self) -> ColorChoice {
142        match &self.inner {
143            StreamInner::PassThrough(_) => ColorChoice::AlwaysAnsi,
144            StreamInner::Strip(_) => ColorChoice::Never,
145            #[cfg(all(windows, feature = "wincon"))]
146            StreamInner::Wincon(_) => ColorChoice::Always,
147        }
148    }
149}
150
151#[cfg(feature = "auto")]
152fn choice(raw: &dyn RawStream) -> ColorChoice {
153    let choice = ColorChoice::global();
154    match choice {
155        ColorChoice::Auto => {
156            let clicolor = anstyle_query::clicolor();
157            let clicolor_enabled = clicolor.unwrap_or(false);
158            let clicolor_disabled = !clicolor.unwrap_or(true);
159            if raw.is_terminal()
160                && !anstyle_query::no_color()
161                && !clicolor_disabled
162                && (anstyle_query::term_supports_color()
163                    || clicolor_enabled
164                    || anstyle_query::is_ci())
165                || anstyle_query::clicolor_force()
166            {
167                ColorChoice::Always
168            } else {
169                ColorChoice::Never
170            }
171        }
172        ColorChoice::AlwaysAnsi | ColorChoice::Always | ColorChoice::Never => choice,
173    }
174}
175
176impl AutoStream<std::io::Stdout> {
177    /// Get exclusive access to the `AutoStream`
178    ///
179    /// Why?
180    /// - Faster performance when writing in a loop
181    /// - Avoid other threads interleaving output with the current thread
182    #[inline]
183    pub fn lock(self) -> AutoStream<std::io::StdoutLock<'static>> {
184        let inner = match self.inner {
185            StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()),
186            StreamInner::Strip(w) => StreamInner::Strip(w.lock()),
187            #[cfg(all(windows, feature = "wincon"))]
188            StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
189        };
190        AutoStream { inner }
191    }
192}
193
194impl AutoStream<std::io::Stderr> {
195    /// Get exclusive access to the `AutoStream`
196    ///
197    /// Why?
198    /// - Faster performance when writing in a loop
199    /// - Avoid other threads interleaving output with the current thread
200    #[inline]
201    pub fn lock(self) -> AutoStream<std::io::StderrLock<'static>> {
202        let inner = match self.inner {
203            StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()),
204            StreamInner::Strip(w) => StreamInner::Strip(w.lock()),
205            #[cfg(all(windows, feature = "wincon"))]
206            StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
207        };
208        AutoStream { inner }
209    }
210}
211
212impl<S> std::io::Write for AutoStream<S>
213where
214    S: RawStream + AsLockedWrite,
215{
216    // Must forward all calls to ensure locking happens appropriately
217    #[inline]
218    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
219        match &mut self.inner {
220            StreamInner::PassThrough(w) => w.as_locked_write().write(buf),
221            StreamInner::Strip(w) => w.write(buf),
222            #[cfg(all(windows, feature = "wincon"))]
223            StreamInner::Wincon(w) => w.write(buf),
224        }
225    }
226    #[inline]
227    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
228        match &mut self.inner {
229            StreamInner::PassThrough(w) => w.as_locked_write().write_vectored(bufs),
230            StreamInner::Strip(w) => w.write_vectored(bufs),
231            #[cfg(all(windows, feature = "wincon"))]
232            StreamInner::Wincon(w) => w.write_vectored(bufs),
233        }
234    }
235    // is_write_vectored: nightly only
236    #[inline]
237    fn flush(&mut self) -> std::io::Result<()> {
238        match &mut self.inner {
239            StreamInner::PassThrough(w) => w.as_locked_write().flush(),
240            StreamInner::Strip(w) => w.flush(),
241            #[cfg(all(windows, feature = "wincon"))]
242            StreamInner::Wincon(w) => w.flush(),
243        }
244    }
245    #[inline]
246    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
247        match &mut self.inner {
248            StreamInner::PassThrough(w) => w.as_locked_write().write_all(buf),
249            StreamInner::Strip(w) => w.write_all(buf),
250            #[cfg(all(windows, feature = "wincon"))]
251            StreamInner::Wincon(w) => w.write_all(buf),
252        }
253    }
254    // write_all_vectored: nightly only
255    #[inline]
256    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
257        match &mut self.inner {
258            StreamInner::PassThrough(w) => w.as_locked_write().write_fmt(args),
259            StreamInner::Strip(w) => w.write_fmt(args),
260            #[cfg(all(windows, feature = "wincon"))]
261            StreamInner::Wincon(w) => w.write_fmt(args),
262        }
263    }
264}