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