foundry_common/io/
shell.rs

1//! Utility functions for writing to [`stdout`](std::io::stdout) and [`stderr`](std::io::stderr).
2//!
3//! Originally from [cargo](https://github.com/rust-lang/cargo/blob/35814255a1dbaeca9219fae81d37a8190050092c/src/cargo/core/shell.rs).
4
5use super::style::*;
6use anstream::AutoStream;
7use anstyle::Style;
8use clap::ValueEnum;
9use eyre::Result;
10use serde::{Deserialize, Serialize};
11use std::{
12    fmt,
13    io::{IsTerminal, prelude::*},
14    ops::DerefMut,
15    sync::{
16        Mutex, OnceLock, PoisonError,
17        atomic::{AtomicBool, Ordering},
18    },
19};
20
21/// Returns the current color choice.
22pub fn color_choice() -> ColorChoice {
23    Shell::get().color_choice()
24}
25
26/// Returns the currently set verbosity level.
27pub fn verbosity() -> Verbosity {
28    Shell::get().verbosity()
29}
30
31/// Set the verbosity level.
32pub fn set_verbosity(verbosity: Verbosity) {
33    Shell::get().set_verbosity(verbosity);
34}
35
36/// Returns whether the output mode is [`OutputMode::Quiet`].
37pub fn is_quiet() -> bool {
38    Shell::get().output_mode().is_quiet()
39}
40
41/// Returns whether stderr is a terminal (tty).
42///
43/// Used to gate progress/spinner output that only makes sense for interactive use.
44pub fn is_err_tty() -> bool {
45    Shell::get().is_err_tty()
46}
47
48/// Returns whether stdout is a terminal (tty).
49///
50/// Used to gate machine-readable stdout records that would duplicate the status prose
51/// already shown on stderr in interactive sessions.
52pub fn is_out_tty() -> bool {
53    Shell::get().is_out_tty()
54}
55
56/// Returns whether the output format is [`OutputFormat::Json`].
57pub fn is_json() -> bool {
58    Shell::get().is_json()
59}
60
61/// Returns whether the output format is [`OutputFormat::Markdown`].
62pub fn is_markdown() -> bool {
63    Shell::get().is_markdown()
64}
65
66/// The global shell instance.
67static GLOBAL_SHELL: OnceLock<Mutex<Shell>> = OnceLock::new();
68
69#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
70/// The requested output mode.
71pub enum OutputMode {
72    /// Default output
73    #[default]
74    Normal,
75    /// No output
76    Quiet,
77}
78
79impl OutputMode {
80    /// Returns true if the output mode is `Normal`.
81    pub fn is_normal(self) -> bool {
82        self == Self::Normal
83    }
84
85    /// Returns true if the output mode is `Quiet`.
86    pub fn is_quiet(self) -> bool {
87        self == Self::Quiet
88    }
89}
90
91/// The requested output format.
92#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
93pub enum OutputFormat {
94    /// Plain text output.
95    #[default]
96    Text,
97    /// JSON output.
98    Json,
99    /// Plain text with markdown tables.
100    Markdown,
101}
102
103impl OutputFormat {
104    /// Returns true if the output format is `Text`.
105    pub fn is_text(self) -> bool {
106        self == Self::Text
107    }
108
109    /// Returns true if the output format is `Json`.
110    pub fn is_json(self) -> bool {
111        self == Self::Json
112    }
113
114    /// Returns true if the output format is `Markdown`.
115    pub fn is_markdown(self) -> bool {
116        self == Self::Markdown
117    }
118}
119
120/// The verbosity level.
121pub type Verbosity = u8;
122
123/// An abstraction around console output that remembers preferences for output
124/// verbosity and color.
125pub struct Shell {
126    /// Wrapper around stdout/stderr. This helps with supporting sending
127    /// output to a memory buffer which is useful for tests.
128    output: ShellOut,
129
130    /// The format to use for message output.
131    output_format: OutputFormat,
132
133    /// The verbosity mode to use for message output.
134    output_mode: OutputMode,
135
136    /// The verbosity level to use for message output.
137    verbosity: Verbosity,
138
139    /// Flag that indicates the current line needs to be cleared before
140    /// printing. Used when a progress bar is currently displayed.
141    needs_clear: AtomicBool,
142}
143
144impl fmt::Debug for Shell {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        let mut s = f.debug_struct("Shell");
147        s.field("output_format", &self.output_format);
148        s.field("output_mode", &self.output_mode);
149        s.field("verbosity", &self.verbosity);
150        if let ShellOut::Stream { color_choice, .. } = self.output {
151            s.field("color_choice", &color_choice);
152        }
153        s.finish()
154    }
155}
156
157/// A `Write`able object, either with or without color support.
158enum ShellOut {
159    /// Color-enabled stdio, with information on whether color should be used.
160    Stream {
161        stdout: AutoStream<std::io::Stdout>,
162        stderr: AutoStream<std::io::Stderr>,
163        stdout_tty: bool,
164        stderr_tty: bool,
165        color_choice: ColorChoice,
166    },
167    /// A write object that ignores all output.
168    Empty(std::io::Empty),
169    /// Captures stdout and stderr into in-memory buffers. Intended for tests.
170    Captured { stdout: Vec<u8>, stderr: Vec<u8> },
171}
172
173/// Whether messages should use color output.
174#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, ValueEnum)]
175pub enum ColorChoice {
176    /// Intelligently guess whether to use color output (default).
177    #[default]
178    Auto,
179    /// Force color output.
180    Always,
181    /// Force disable color output.
182    Never,
183}
184
185impl Default for Shell {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191impl Shell {
192    /// Creates a new shell (color choice and verbosity), defaulting to 'auto' color and verbose
193    /// output.
194    pub fn new() -> Self {
195        Self::new_with(
196            OutputFormat::Text,
197            OutputMode::Normal,
198            ColorChoice::Auto,
199            Verbosity::default(),
200        )
201    }
202
203    /// Creates a new shell with the given color choice and verbosity.
204    pub fn new_with(
205        format: OutputFormat,
206        mode: OutputMode,
207        color: ColorChoice,
208        verbosity: Verbosity,
209    ) -> Self {
210        Self {
211            output: ShellOut::Stream {
212                stdout: AutoStream::new(std::io::stdout(), color.to_anstream_color_choice()),
213                stderr: AutoStream::new(std::io::stderr(), color.to_anstream_color_choice()),
214                color_choice: color,
215                stdout_tty: std::io::stdout().is_terminal(),
216                stderr_tty: std::io::stderr().is_terminal(),
217            },
218            output_format: format,
219            output_mode: mode,
220            verbosity,
221            needs_clear: AtomicBool::new(false),
222        }
223    }
224
225    /// Creates a shell that ignores all output.
226    pub const fn empty() -> Self {
227        Self {
228            output: ShellOut::Empty(std::io::empty()),
229            output_format: OutputFormat::Text,
230            output_mode: OutputMode::Quiet,
231            verbosity: 0,
232            needs_clear: AtomicBool::new(false),
233        }
234    }
235
236    /// Creates a shell that captures stdout and stderr into in-memory buffers.
237    ///
238    /// Intended for tests that want to assert how a piece of code routes output
239    /// between stdout and stderr. Use [`Shell::captured_stdout`] and
240    /// [`Shell::captured_stderr`] to read the buffers back.
241    pub const fn captured() -> Self {
242        Self {
243            output: ShellOut::Captured { stdout: Vec::new(), stderr: Vec::new() },
244            output_format: OutputFormat::Text,
245            output_mode: OutputMode::Normal,
246            verbosity: 0,
247            needs_clear: AtomicBool::new(false),
248        }
249    }
250
251    /// Returns the captured stdout buffer, if this shell was created via [`Shell::captured`].
252    pub fn captured_stdout(&self) -> Option<&[u8]> {
253        match &self.output {
254            ShellOut::Captured { stdout, .. } => Some(stdout),
255            _ => None,
256        }
257    }
258
259    /// Returns the captured stderr buffer, if this shell was created via [`Shell::captured`].
260    pub fn captured_stderr(&self) -> Option<&[u8]> {
261        match &self.output {
262            ShellOut::Captured { stderr, .. } => Some(stderr),
263            _ => None,
264        }
265    }
266
267    /// Acquire a lock to the global shell.
268    ///
269    /// Initializes it with the default values if it has not been set yet.
270    pub fn get() -> impl DerefMut<Target = Self> + 'static {
271        GLOBAL_SHELL.get_or_init(Default::default).lock().unwrap_or_else(PoisonError::into_inner)
272    }
273
274    /// Set the global shell.
275    ///
276    /// # Panics
277    ///
278    /// Panics if the global shell has already been set.
279    #[track_caller]
280    pub fn set(self) {
281        GLOBAL_SHELL
282            .set(Mutex::new(self))
283            .unwrap_or_else(|_| panic!("attempted to set global shell twice"))
284    }
285
286    /// Sets whether the next print should clear the current line and returns the previous value.
287    pub fn set_needs_clear(&self, needs_clear: bool) -> bool {
288        self.needs_clear.swap(needs_clear, Ordering::Relaxed)
289    }
290
291    /// Returns `true` if the output format is JSON.
292    pub fn is_json(&self) -> bool {
293        self.output_format.is_json()
294    }
295
296    /// Returns `true` if the output format is Markdown.
297    pub fn is_markdown(&self) -> bool {
298        self.output_format.is_markdown()
299    }
300
301    /// Returns `true` if the verbosity level is `Quiet`.
302    pub fn is_quiet(&self) -> bool {
303        self.output_mode.is_quiet()
304    }
305
306    /// Returns `true` if the `needs_clear` flag is set.
307    pub fn needs_clear(&self) -> bool {
308        self.needs_clear.load(Ordering::Relaxed)
309    }
310
311    /// Gets the output format of the shell.
312    pub const fn output_format(&self) -> OutputFormat {
313        self.output_format
314    }
315
316    /// Gets the output mode of the shell.
317    pub const fn output_mode(&self) -> OutputMode {
318        self.output_mode
319    }
320
321    /// Gets the verbosity of the shell when [`OutputMode::Normal`] is set.
322    pub const fn verbosity(&self) -> Verbosity {
323        self.verbosity
324    }
325
326    /// Sets the verbosity level.
327    pub const fn set_verbosity(&mut self, verbosity: Verbosity) {
328        self.verbosity = verbosity;
329    }
330
331    /// Sets the output mode.
332    pub const fn set_output_mode(&mut self, output_mode: OutputMode) {
333        self.output_mode = output_mode;
334    }
335
336    /// Gets the current color choice.
337    ///
338    /// If we are not using a color stream, this will always return `Never`, even if the color
339    /// choice has been set to something else.
340    pub const fn color_choice(&self) -> ColorChoice {
341        match self.output {
342            ShellOut::Stream { color_choice, .. } => color_choice,
343            ShellOut::Empty(_) | ShellOut::Captured { .. } => ColorChoice::Never,
344        }
345    }
346
347    /// Returns `true` if stderr is a tty.
348    pub const fn is_err_tty(&self) -> bool {
349        match self.output {
350            ShellOut::Stream { stderr_tty, .. } => stderr_tty,
351            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
352        }
353    }
354
355    /// Returns `true` if stdout is a tty.
356    pub const fn is_out_tty(&self) -> bool {
357        match self.output {
358            ShellOut::Stream { stdout_tty, .. } => stdout_tty,
359            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
360        }
361    }
362
363    /// Whether `stderr` supports color.
364    pub fn err_supports_color(&self) -> bool {
365        match &self.output {
366            ShellOut::Stream { stderr, .. } => supports_color(stderr.current_choice()),
367            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
368        }
369    }
370
371    /// Whether `stdout` supports color.
372    pub fn out_supports_color(&self) -> bool {
373        match &self.output {
374            ShellOut::Stream { stdout, .. } => supports_color(stdout.current_choice()),
375            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
376        }
377    }
378
379    /// Gets a reference to the underlying stdout writer.
380    pub fn out(&mut self) -> &mut dyn Write {
381        self.maybe_err_erase_line();
382        self.output.stdout()
383    }
384
385    /// Gets a reference to the underlying stderr writer.
386    pub fn err(&mut self) -> &mut dyn Write {
387        self.maybe_err_erase_line();
388        self.output.stderr()
389    }
390
391    /// Erase from cursor to end of line if needed.
392    pub fn maybe_err_erase_line(&mut self) {
393        if self.err_supports_color() && self.set_needs_clear(false) {
394            // This is the "EL - Erase in Line" sequence. It clears from the cursor
395            // to the end of line.
396            // https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
397            let _ = self.output.stderr().write_all(b"\x1B[K");
398        }
399    }
400
401    /// Prints a red 'error' message. Use the [`sh_err!`] macro instead.
402    /// This will render a message in [ERROR] style with a bold `Error: ` prefix.
403    ///
404    /// **Note**: will log regardless of the verbosity level.
405    pub fn error(&mut self, message: impl fmt::Display) -> Result<()> {
406        self.maybe_err_erase_line();
407        self.output.message_stderr(&"Error", &ERROR, Some(&message), false)
408    }
409
410    /// Prints an amber 'warning' message. Use the [`sh_warn!`] macro instead.
411    /// This will render a message in [WARN] style with a bold `Warning: `prefix.
412    ///
413    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
414    pub fn warn(&mut self, message: impl fmt::Display) -> Result<()> {
415        match self.output_mode {
416            OutputMode::Quiet => Ok(()),
417            _ => self.print(&"Warning", &WARN, Some(&message), false),
418        }
419    }
420
421    /// Write a styled fragment.
422    ///
423    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
424    pub fn write_stdout(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
425        self.output.write_stdout(fragment, color)
426    }
427
428    /// Write a styled fragment with the default color. Use the [`sh_print!`] macro instead.
429    ///
430    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
431    //
432    // TODO: stdout is the canonical machine-readable result of a command and should NOT be
433    // suppressed by `--quiet` (see `docs/dev/output-channels.md`). Flip this once the major
434    // prose `sh_println!` call sites in forge/script have been migrated to `sh_status!`.
435    pub fn print_out(&mut self, fragment: impl fmt::Display) -> Result<()> {
436        match self.output_mode {
437            OutputMode::Quiet => Ok(()),
438            _ => self.write_stdout(fragment, &Style::new()),
439        }
440    }
441
442    /// Write a styled fragment
443    ///
444    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
445    pub fn write_stderr(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
446        self.output.write_stderr(fragment, color)
447    }
448
449    /// Write a styled fragment with the default color. Use the [`sh_eprint!`] macro instead.
450    ///
451    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
452    pub fn print_err(&mut self, fragment: impl fmt::Display) -> Result<()> {
453        match self.output_mode {
454            OutputMode::Quiet => Ok(()),
455            _ => self.write_stderr(fragment, &Style::new()),
456        }
457    }
458
459    /// Prints a message, where the status will have `color` color, and can be justified. The
460    /// messages follows without color.
461    fn print(
462        &mut self,
463        status: &dyn fmt::Display,
464        style: &Style,
465        message: Option<&dyn fmt::Display>,
466        justified: bool,
467    ) -> Result<()> {
468        match self.output_mode {
469            OutputMode::Quiet => Ok(()),
470            _ => {
471                self.maybe_err_erase_line();
472                self.output.message_stderr(status, style, message, justified)
473            }
474        }
475    }
476}
477
478impl ShellOut {
479    /// Prints out a message with a status to stderr. The status comes first, and is bold plus the
480    /// given color. The status can be justified, in which case the max width that will right
481    /// align is 12 chars.
482    fn message_stderr(
483        &mut self,
484        status: &dyn fmt::Display,
485        style: &Style,
486        message: Option<&dyn fmt::Display>,
487        justified: bool,
488    ) -> Result<()> {
489        let buffer = Self::format_message(status, message, style, justified)?;
490        self.stderr().write_all(&buffer)?;
491        Ok(())
492    }
493
494    /// Write a styled fragment
495    fn write_stdout(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
496        let mut buffer = Vec::new();
497        write!(buffer, "{style}{fragment}{style:#}")?;
498        self.stdout().write_all(&buffer)?;
499        Ok(())
500    }
501
502    /// Write a styled fragment
503    fn write_stderr(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
504        let mut buffer = Vec::new();
505        write!(buffer, "{style}{fragment}{style:#}")?;
506        self.stderr().write_all(&buffer)?;
507        Ok(())
508    }
509
510    /// Gets stdout as a [`io::Write`](Write) trait object.
511    fn stdout(&mut self) -> &mut dyn Write {
512        match self {
513            Self::Stream { stdout, .. } => stdout,
514            Self::Empty(e) => e,
515            Self::Captured { stdout, .. } => stdout,
516        }
517    }
518
519    /// Gets stderr as a [`io::Write`](Write) trait object.
520    fn stderr(&mut self) -> &mut dyn Write {
521        match self {
522            Self::Stream { stderr, .. } => stderr,
523            Self::Empty(e) => e,
524            Self::Captured { stderr, .. } => stderr,
525        }
526    }
527
528    /// Formats a message with a status and optional message.
529    fn format_message(
530        status: &dyn fmt::Display,
531        message: Option<&dyn fmt::Display>,
532        style: &Style,
533        justified: bool,
534    ) -> Result<Vec<u8>> {
535        let bold = anstyle::Style::new().bold();
536
537        let mut buffer = Vec::new();
538        if justified {
539            write!(buffer, "{style}{status:>12}{style:#}")?;
540        } else {
541            write!(buffer, "{style}{status}{style:#}{bold}:{bold:#}")?;
542        }
543        match message {
544            Some(message) => {
545                writeln!(buffer, " {message}")?;
546            }
547            None => write!(buffer, " ")?,
548        }
549
550        Ok(buffer)
551    }
552}
553
554impl ColorChoice {
555    /// Converts our color choice to [`anstream`]'s version.
556    const fn to_anstream_color_choice(self) -> anstream::ColorChoice {
557        match self {
558            Self::Always => anstream::ColorChoice::Always,
559            Self::Never => anstream::ColorChoice::Never,
560            Self::Auto => anstream::ColorChoice::Auto,
561        }
562    }
563}
564
565const fn supports_color(choice: anstream::ColorChoice) -> bool {
566    match choice {
567        anstream::ColorChoice::Always
568        | anstream::ColorChoice::AlwaysAnsi
569        | anstream::ColorChoice::Auto => true,
570        anstream::ColorChoice::Never => false,
571    }
572}
foundry_common/io/shell.rs

foundry_common/io/
shell.rs
