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 #[inline]
161 pub fn is_terminal(&self) -> bool {
162 match &self.inner {
163 StreamInner::PassThrough(w) => w.is_terminal(),
164 StreamInner::Strip(w) => w.is_terminal(),
165 #[cfg(all(windows, feature = "wincon"))]
166 StreamInner::Wincon(_) => true, // its only ever a terminal
167 }
168 }
169
170 /// Prefer [`AutoStream::choice`]
171 ///
172 /// This doesn't report what is requested but what is currently active.
173 #[inline]
174 #[cfg(feature = "auto")]
175 pub fn current_choice(&self) -> ColorChoice {
176 match &self.inner {
177 StreamInner::PassThrough(_) => ColorChoice::AlwaysAnsi,
178 StreamInner::Strip(_) => ColorChoice::Never,
179 #[cfg(all(windows, feature = "wincon"))]
180 StreamInner::Wincon(_) => ColorChoice::Always,
181 }
182 }
183}
184
185#[cfg(feature = "auto")]
186fn choice(raw: &dyn RawStream) -> ColorChoice {
187 let choice = ColorChoice::global();
188 match choice {
189 ColorChoice::Auto => {
190 let clicolor = anstyle_query::clicolor();
191 let clicolor_enabled = clicolor.unwrap_or(false);
192 let clicolor_disabled = !clicolor.unwrap_or(true);
193 if anstyle_query::no_color() {
194 ColorChoice::Never
195 } else if anstyle_query::clicolor_force() {
196 ColorChoice::Always
197 } else if clicolor_disabled {
198 ColorChoice::Never
199 } else if raw.is_terminal()
200 && (anstyle_query::term_supports_color()
201 || clicolor_enabled
202 || anstyle_query::is_ci())
203 {
204 ColorChoice::Always
205 } else {
206 ColorChoice::Never
207 }
208 }
209 ColorChoice::AlwaysAnsi | ColorChoice::Always | ColorChoice::Never => choice,
210 }
211}
212
213impl AutoStream<std::io::Stdout> {
214 /// Get exclusive access to the `AutoStream`
215 ///
216 /// Why?
217 /// - Faster performance when writing in a loop
218 /// - Avoid other threads interleaving output with the current thread
219 #[inline]
220 pub fn lock(self) -> AutoStream<std::io::StdoutLock<'static>> {
221 let inner: StreamInner> = match self.inner {
222 StreamInner::PassThrough(w: Stdout) => StreamInner::PassThrough(w.lock()),
223 StreamInner::Strip(w: StripStream) => StreamInner::Strip(w.lock()),
224 #[cfg(all(windows, feature = "wincon"))]
225 StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
226 };
227 AutoStream { inner }
228 }
229}
230
231impl AutoStream<std::io::Stderr> {
232 /// Get exclusive access to the `AutoStream`
233 ///
234 /// Why?
235 /// - Faster performance when writing in a loop
236 /// - Avoid other threads interleaving output with the current thread
237 #[inline]
238 pub fn lock(self) -> AutoStream<std::io::StderrLock<'static>> {
239 let inner: StreamInner> = match self.inner {
240 StreamInner::PassThrough(w: Stderr) => StreamInner::PassThrough(w.lock()),
241 StreamInner::Strip(w: StripStream) => StreamInner::Strip(w.lock()),
242 #[cfg(all(windows, feature = "wincon"))]
243 StreamInner::Wincon(w) => StreamInner::Wincon(w.lock()),
244 };
245 AutoStream { inner }
246 }
247}
248
249impl<S> std::io::Write for AutoStream<S>
250where
251 S: RawStream + AsLockedWrite,
252{
253 // Must forward all calls to ensure locking happens appropriately
254 #[inline]
255 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
256 match &mut self.inner {
257 StreamInner::PassThrough(w) => w.as_locked_write().write(buf),
258 StreamInner::Strip(w) => w.write(buf),
259 #[cfg(all(windows, feature = "wincon"))]
260 StreamInner::Wincon(w) => w.write(buf),
261 }
262 }
263 #[inline]
264 fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
265 match &mut self.inner {
266 StreamInner::PassThrough(w) => w.as_locked_write().write_vectored(bufs),
267 StreamInner::Strip(w) => w.write_vectored(bufs),
268 #[cfg(all(windows, feature = "wincon"))]
269 StreamInner::Wincon(w) => w.write_vectored(bufs),
270 }
271 }
272 // is_write_vectored: nightly only
273 #[inline]
274 fn flush(&mut self) -> std::io::Result<()> {
275 match &mut self.inner {
276 StreamInner::PassThrough(w) => w.as_locked_write().flush(),
277 StreamInner::Strip(w) => w.flush(),
278 #[cfg(all(windows, feature = "wincon"))]
279 StreamInner::Wincon(w) => w.flush(),
280 }
281 }
282 #[inline]
283 fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
284 match &mut self.inner {
285 StreamInner::PassThrough(w) => w.as_locked_write().write_all(buf),
286 StreamInner::Strip(w) => w.write_all(buf),
287 #[cfg(all(windows, feature = "wincon"))]
288 StreamInner::Wincon(w) => w.write_all(buf),
289 }
290 }
291 // write_all_vectored: nightly only
292 #[inline]
293 fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
294 match &mut self.inner {
295 StreamInner::PassThrough(w) => w.as_locked_write().write_fmt(args),
296 StreamInner::Strip(w) => w.write_fmt(args),
297 #[cfg(all(windows, feature = "wincon"))]
298 StreamInner::Wincon(w) => w.write_fmt(args),
299 }
300 }
301}
302