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///
10/// This includes
11/// - Stripping colors for non-terminals
12/// - Respecting env variables like [NO_COLOR](https://no-color.org/) or [CLICOLOR](https://bixense.com/clicolors/)
13/// - *(windows)* Falling back to the wincon API where [ENABLE_VIRTUAL_TERMINAL_PROCESSING](https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences#output-sequences) is unsupported
14///
15/// You can customize auto-detection by calling into
16/// [anstyle_query](https://docs.rs/anstyle-query/latest/anstyle_query/)
17/// to get a [`ColorChoice`] and then calling [`AutoStream::new(stream, choice)`].
18#[derive(Debug)]
19pub struct AutoStream<S: RawStream> {
20    inner: StreamInner<S>,
21}
22
23#[derive(Debug)]
24enum StreamInner<S: RawStream> {
25    PassThrough(S),
26    Strip(StripStream<S>),
27    #[cfg(all(windows, feature = "wincon"))]
28    Wincon(WinconStream<S>),
29}
30
31impl<S> AutoStream<S>
32where
33    S: RawStream,
34{
35    /// Runtime control over styling behavior
36    ///
37    /// # Example
38    ///
39    /// ```rust
40    /// # #[cfg(feature = "auto")] {
41    /// # use std::io::IsTerminal as _;
42    /// // Like `AutoStream::choice` but without `NO_COLOR`, `CLICOLOR_FORCE`, `CI`
43    /// fn choice(raw: &dyn anstream::stream::RawStream) -> anstream::ColorChoice {
44    ///     let choice = anstream::ColorChoice::global();
45    ///     if choice == anstream::ColorChoice::Auto {
46    ///         if raw.is_terminal() && anstyle_query::term_supports_color() {
47    ///             anstream::ColorChoice::Always
48    ///         } else {
49    ///             anstream::ColorChoice::Never
50    ///         }
51    ///     } else {
52    ///         choice
53    ///     }
54    /// }
55    ///
56    /// let stream = std::io::stdout();
57    /// let choice = choice(&stream);
58    /// let auto = anstream::AutoStream::new(stream, choice);
59    /// # }
60    /// ```
61    #[inline]
62    pub fn new(raw: S, choice: ColorChoice) -> Self {
63        match choice {
64            #[cfg(feature = "auto")]
65            ColorChoice::Auto => Self::auto(raw),
66            #[cfg(not(feature = "auto"))]
67            ColorChoice::Auto => Self::never(raw),
68            ColorChoice::AlwaysAnsi => Self::always_ansi(raw),
69            ColorChoice::Always => Self::always(raw),
70            ColorChoice::Never => Self::never(raw),
71        }
72    }
73
74    /// Auto-adapt for the stream's capabilities
75    #[cfg(feature = "auto")]
76    #[inline]
77    pub fn auto(raw: S) -> Self {
78        let choice = Self::choice(&raw);
79        debug_assert_ne!(choice, ColorChoice::Auto);
80        Self::new(raw, choice)
81    }
82
83    /// Report the desired choice for the given stream
84    #[cfg(feature = "auto")]
85    pub fn choice(raw: &S) -> ColorChoice {
86        choice(raw)
87    }
88
89    /// Force ANSI escape codes to be passed through as-is, no matter what the inner `Write`
90    /// supports.
91    #[inline]
92    pub fn always_ansi(raw: S) -> Self {
93        #[cfg(feature = "auto")]
94        {
95            if raw.is_terminal() {
96                let _ = anstyle_query::windows::enable_ansi_colors();
97            }
98        }
99        Self::always_ansi_(raw)
100    }
101
102    #[inline]
103    fn always_ansi_(raw: S) -> Self {
104        let inner = StreamInner::PassThrough(raw);
105        AutoStream { inner }
106    }
107
108    /// Force color, no matter what the inner `Write` supports.
109    #[inline]
110    pub fn always(raw: S) -> Self {
111        if cfg!(windows) {
112            #[cfg(feature = "auto")]
113            let use_wincon = raw.is_terminal()
114                && !anstyle_query::windows::enable_ansi_colors().unwrap_or(true)
115                && !anstyle_query::term_supports_ansi_color();
116            #[cfg(not(feature = "auto"))]
117            let use_wincon = true;
118            if use_wincon {
119                Self::wincon(raw).unwrap_or_else(|raw| Self::always_ansi_(raw))
120            } else {
121                Self::always_ansi_(raw)
122            }
123        } else {
124            Self::always_ansi(raw)
125        }
126    }
127
128    /// Only pass printable data to the inner `Write`.
129    #[inline]
130    pub fn never(raw: S) -> Self {
131        let inner = StreamInner::Strip(StripStream::new(raw));
132        AutoStream { inner }
133    }
134
135    #[inline]
136    fn wincon(raw: S) -> Result<Self, S> {
137        #[cfg(all(windows, feature = "wincon"))]
138        {
139            Ok(Self {
140                inner: StreamInner::Wincon(WinconStream::new(raw)),
141            })
142        }
143        #[cfg(not(all(windows, feature = "wincon")))]
144        {
145            Err(raw)
146        }
147    }
148
149    /// Get the wrapped [`RawStream`]
150    #[inline]
151    pub fn into_inner(self) -> S {
152        match self.inner {
153            StreamInner::PassThrough(w) => w,
154            StreamInner::Strip(w) => w.into_inner(),
155            #[cfg(all(windows, feature = "wincon"))]
156            StreamInner::Wincon(w) => w.into_inner(),
157        }
158    }
159
160    /// Get the wrapped [`RawStream`]
161    #[inline]
162    pub fn as_inner(&self) -> &S {
163        match &self.inner {
164            StreamInner::PassThrough(w) => w,
165            StreamInner::Strip(w) => w.as_inner(),
166            #[cfg(all(windows, feature = "wincon"))]
167            StreamInner::Wincon(w) => w.as_inner(),
168        }
169    }
170
171    /// Returns `true` if the descriptor/handle refers to a terminal/tty.
172    #[inline]
173    pub fn is_terminal(&self) -> bool {
174        match &self.inner {
175            StreamInner::PassThrough(w) => w.is_terminal(),
176            StreamInner::Strip(w) => w.is_terminal(),
177            #[cfg(all(windows, feature = "wincon"))]
178            StreamInner::Wincon(_) => true, // its only ever a terminal
179        }
180    }
181
182    /// Prefer [`AutoStream::choice`]
183    ///
184    /// This doesn't report what is requested but what is currently active.
185    #[inline]
186    #[cfg(feature = "auto")]
187    pub fn current_choice(&self) -> ColorChoice {
188        match &self.inner {
189            StreamInner::PassThrough(_) => ColorChoice::AlwaysAnsi,
190            StreamInner::Strip(_) => ColorChoice::Never,
191            #[cfg(all(windows, feature = "wincon"))]
192            StreamInner::Wincon(_) => ColorChoice::Always,
193        }
194    }
195}
196
197#[cfg(feature = "auto")]
198fn choice(raw: &dyn RawStream) -> ColorChoice {
199    let choice = ColorChoice::global();
200    match choice {
201        ColorChoice::Auto => {
202            let clicolor = anstyle_query::clicolor();
203            let clicolor_enabled = clicolor.unwrap_or(false);
204            let clicolor_disabled = !clicolor.unwrap_or(true);
205            if anstyle_query::no_color() {
206                ColorChoice::Never
207            } else if anstyle_query::clicolor_force() {
208                ColorChoice::Always
209            } else if clicolor_disabled {
210                ColorChoice::Never
211            } else if raw.is_terminal()
212                && (anstyle_query::term_supports_color()
213                    || clicolor_enabled
214                    || anstyle_query::is_ci())
215            {
216                ColorChoice::Always
217            } else {
218                ColorChoice::Never
219            }
220        }
221        ColorChoice::AlwaysAnsi | ColorChoice::Always | ColorChoice::Never => choice,
222    }
223}
224
225impl AutoStream<std::io::Stdout> {
226    /// Get exclusive access to the `AutoStream`
227    ///
228    /// Why?
229    /// - Faster performance when writing in a loop
230    /// - Avoid other threads interleaving output with the current thread
231    #[inline]
232    pub fn lock(self) -> AutoStream<std::io::StdoutLock<'static>> {
233        let inner = match self.inner {
234            StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()),
235            StreamInner::Strip(w) => StreamInner::Strip(w.lock()),
236            #[cfg(all(windows, feature = "wincon"))]
237            StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
238        };
239        AutoStream { inner }
240    }
241}
242
243impl AutoStream<std::io::Stderr> {
244    /// Get exclusive access to the `AutoStream`
245    ///
246    /// Why?
247    /// - Faster performance when writing in a loop
248    /// - Avoid other threads interleaving output with the current thread
249    #[inline]
250    pub fn lock(self) -> AutoStream<std::io::StderrLock<'static>> {
251        let inner = match self.inner {
252            StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()),
253            StreamInner::Strip(w) => StreamInner::Strip(w.lock()),
254            #[cfg(all(windows, feature = "wincon"))]
255            StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
256        };
257        AutoStream { inner }
258    }
259}
260
261impl<S> std::io::Write for AutoStream<S>
262where
263    S: RawStream + AsLockedWrite,
264{
265    // Must forward all calls to ensure locking happens appropriately
266    #[inline]
267    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
268        match &mut self.inner {
269            StreamInner::PassThrough(w) => w.as_locked_write().write(buf),
270            StreamInner::Strip(w) => w.write(buf),
271            #[cfg(all(windows, feature = "wincon"))]
272            StreamInner::Wincon(w) => w.write(buf),
273        }
274    }
275    #[inline]
276    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
277        match &mut self.inner {
278            StreamInner::PassThrough(w) => w.as_locked_write().write_vectored(bufs),
279            StreamInner::Strip(w) => w.write_vectored(bufs),
280            #[cfg(all(windows, feature = "wincon"))]
281            StreamInner::Wincon(w) => w.write_vectored(bufs),
282        }
283    }
284    // is_write_vectored: nightly only
285    #[inline]
286    fn flush(&mut self) -> std::io::Result<()> {
287        match &mut self.inner {
288            StreamInner::PassThrough(w) => w.as_locked_write().flush(),
289            StreamInner::Strip(w) => w.flush(),
290            #[cfg(all(windows, feature = "wincon"))]
291            StreamInner::Wincon(w) => w.flush(),
292        }
293    }
294    #[inline]
295    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
296        match &mut self.inner {
297            StreamInner::PassThrough(w) => w.as_locked_write().write_all(buf),
298            StreamInner::Strip(w) => w.write_all(buf),
299            #[cfg(all(windows, feature = "wincon"))]
300            StreamInner::Wincon(w) => w.write_all(buf),
301        }
302    }
303    // write_all_vectored: nightly only
304    #[inline]
305    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
306        match &mut self.inner {
307            StreamInner::PassThrough(w) => w.as_locked_write().write_fmt(args),
308            StreamInner::Strip(w) => w.write_fmt(args),
309            #[cfg(all(windows, feature = "wincon"))]
310            StreamInner::Wincon(w) => w.write_fmt(args),
311        }
312    }
313}
anstream/auto.rs

anstream/
auto.rs
