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    /// Returns `true` if the `needs_clear` flag is unset.
312    pub fn is_cleared(&self) -> bool {
313        !self.needs_clear()
314    }
315
316    /// Gets the output format of the shell.
317    pub const fn output_format(&self) -> OutputFormat {
318        self.output_format
319    }
320
321    /// Gets the output mode of the shell.
322    pub const fn output_mode(&self) -> OutputMode {
323        self.output_mode
324    }
325
326    /// Gets the verbosity of the shell when [`OutputMode::Normal`] is set.
327    pub const fn verbosity(&self) -> Verbosity {
328        self.verbosity
329    }
330
331    /// Sets the verbosity level.
332    pub const fn set_verbosity(&mut self, verbosity: Verbosity) {
333        self.verbosity = verbosity;
334    }
335
336    /// Sets the output mode.
337    pub const fn set_output_mode(&mut self, output_mode: OutputMode) {
338        self.output_mode = output_mode;
339    }
340
341    /// Gets the current color choice.
342    ///
343    /// If we are not using a color stream, this will always return `Never`, even if the color
344    /// choice has been set to something else.
345    pub const fn color_choice(&self) -> ColorChoice {
346        match self.output {
347            ShellOut::Stream { color_choice, .. } => color_choice,
348            ShellOut::Empty(_) | ShellOut::Captured { .. } => ColorChoice::Never,
349        }
350    }
351
352    /// Returns `true` if stderr is a tty.
353    pub const fn is_err_tty(&self) -> bool {
354        match self.output {
355            ShellOut::Stream { stderr_tty, .. } => stderr_tty,
356            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
357        }
358    }
359
360    /// Returns `true` if stdout is a tty.
361    pub const fn is_out_tty(&self) -> bool {
362        match self.output {
363            ShellOut::Stream { stdout_tty, .. } => stdout_tty,
364            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
365        }
366    }
367
368    /// Whether `stderr` supports color.
369    pub fn err_supports_color(&self) -> bool {
370        match &self.output {
371            ShellOut::Stream { stderr, .. } => supports_color(stderr.current_choice()),
372            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
373        }
374    }
375
376    /// Whether `stdout` supports color.
377    pub fn out_supports_color(&self) -> bool {
378        match &self.output {
379            ShellOut::Stream { stdout, .. } => supports_color(stdout.current_choice()),
380            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
381        }
382    }
383
384    /// Gets a reference to the underlying stdout writer.
385    pub fn out(&mut self) -> &mut dyn Write {
386        self.maybe_err_erase_line();
387        self.output.stdout()
388    }
389
390    /// Gets a reference to the underlying stderr writer.
391    pub fn err(&mut self) -> &mut dyn Write {
392        self.maybe_err_erase_line();
393        self.output.stderr()
394    }
395
396    /// Erase from cursor to end of line if needed.
397    pub fn maybe_err_erase_line(&mut self) {
398        if self.err_supports_color() && self.set_needs_clear(false) {
399            // This is the "EL - Erase in Line" sequence. It clears from the cursor
400            // to the end of line.
401            // https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
402            let _ = self.output.stderr().write_all(b"\x1B[K");
403        }
404    }
405
406    /// Prints a red 'error' message. Use the [`sh_err!`] macro instead.
407    /// This will render a message in [ERROR] style with a bold `Error: ` prefix.
408    ///
409    /// **Note**: will log regardless of the verbosity level.
410    pub fn error(&mut self, message: impl fmt::Display) -> Result<()> {
411        self.maybe_err_erase_line();
412        self.output.message_stderr(&"Error", &ERROR, Some(&message), false)
413    }
414
415    /// Prints an amber 'warning' message. Use the [`sh_warn!`] macro instead.
416    /// This will render a message in [WARN] style with a bold `Warning: `prefix.
417    ///
418    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
419    pub fn warn(&mut self, message: impl fmt::Display) -> Result<()> {
420        match self.output_mode {
421            OutputMode::Quiet => Ok(()),
422            _ => self.print(&"Warning", &WARN, Some(&message), false),
423        }
424    }
425
426    /// Write a styled fragment.
427    ///
428    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
429    pub fn write_stdout(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
430        self.output.write_stdout(fragment, color)
431    }
432
433    /// Write a styled fragment with the default color. Use the [`sh_print!`] macro instead.
434    ///
435    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
436    //
437    // TODO: stdout is the canonical machine-readable result of a command and should NOT be
438    // suppressed by `--quiet` (see `docs/dev/output-channels.md`). Flip this once the major
439    // prose `sh_println!` call sites in forge/script have been migrated to `sh_status!`.
440    pub fn print_out(&mut self, fragment: impl fmt::Display) -> Result<()> {
441        match self.output_mode {
442            OutputMode::Quiet => Ok(()),
443            _ => self.write_stdout(fragment, &Style::new()),
444        }
445    }
446
447    /// Write a styled fragment
448    ///
449    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
450    pub fn write_stderr(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
451        self.output.write_stderr(fragment, color)
452    }
453
454    /// Write a styled fragment with the default color. Use the [`sh_eprint!`] macro instead.
455    ///
456    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
457    pub fn print_err(&mut self, fragment: impl fmt::Display) -> Result<()> {
458        match self.output_mode {
459            OutputMode::Quiet => Ok(()),
460            _ => self.write_stderr(fragment, &Style::new()),
461        }
462    }
463
464    /// Prints a message, where the status will have `color` color, and can be justified. The
465    /// messages follows without color.
466    fn print(
467        &mut self,
468        status: &dyn fmt::Display,
469        style: &Style,
470        message: Option<&dyn fmt::Display>,
471        justified: bool,
472    ) -> Result<()> {
473        match self.output_mode {
474            OutputMode::Quiet => Ok(()),
475            _ => {
476                self.maybe_err_erase_line();
477                self.output.message_stderr(status, style, message, justified)
478            }
479        }
480    }
481}
482
483impl ShellOut {
484    /// Prints out a message with a status to stderr. The status comes first, and is bold plus the
485    /// given color. The status can be justified, in which case the max width that will right
486    /// align is 12 chars.
487    fn message_stderr(
488        &mut self,
489        status: &dyn fmt::Display,
490        style: &Style,
491        message: Option<&dyn fmt::Display>,
492        justified: bool,
493    ) -> Result<()> {
494        let buffer = Self::format_message(status, message, style, justified)?;
495        self.stderr().write_all(&buffer)?;
496        Ok(())
497    }
498
499    /// Write a styled fragment
500    fn write_stdout(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
501        let mut buffer = Vec::new();
502        write!(buffer, "{style}{fragment}{style:#}")?;
503        self.stdout().write_all(&buffer)?;
504        Ok(())
505    }
506
507    /// Write a styled fragment
508    fn write_stderr(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
509        let mut buffer = Vec::new();
510        write!(buffer, "{style}{fragment}{style:#}")?;
511        self.stderr().write_all(&buffer)?;
512        Ok(())
513    }
514
515    /// Gets stdout as a [`io::Write`](Write) trait object.
516    fn stdout(&mut self) -> &mut dyn Write {
517        match self {
518            Self::Stream { stdout, .. } => stdout,
519            Self::Empty(e) => e,
520            Self::Captured { stdout, .. } => stdout,
521        }
522    }
523
524    /// Gets stderr as a [`io::Write`](Write) trait object.
525    fn stderr(&mut self) -> &mut dyn Write {
526        match self {
527            Self::Stream { stderr, .. } => stderr,
528            Self::Empty(e) => e,
529            Self::Captured { stderr, .. } => stderr,
530        }
531    }
532
533    /// Formats a message with a status and optional message.
534    fn format_message(
535        status: &dyn fmt::Display,
536        message: Option<&dyn fmt::Display>,
537        style: &Style,
538        justified: bool,
539    ) -> Result<Vec<u8>> {
540        let bold = anstyle::Style::new().bold();
541
542        let mut buffer = Vec::new();
543        if justified {
544            write!(buffer, "{style}{status:>12}{style:#}")?;
545        } else {
546            write!(buffer, "{style}{status}{style:#}{bold}:{bold:#}")?;
547        }
548        match message {
549            Some(message) => {
550                writeln!(buffer, " {message}")?;
551            }
552            None => write!(buffer, " ")?,
553        }
554
555        Ok(buffer)
556    }
557}
558
559impl ColorChoice {
560    /// Converts our color choice to [`anstream`]'s version.
561    const fn to_anstream_color_choice(self) -> anstream::ColorChoice {
562        match self {
563            Self::Always => anstream::ColorChoice::Always,
564            Self::Never => anstream::ColorChoice::Never,
565            Self::Auto => anstream::ColorChoice::Auto,
566        }
567    }
568}
569
570const fn supports_color(choice: anstream::ColorChoice) -> bool {
571    match choice {
572        anstream::ColorChoice::Always
573        | anstream::ColorChoice::AlwaysAnsi
574        | anstream::ColorChoice::Auto => true,
575        anstream::ColorChoice::Never => false,
576    }
577}
foundry_common/io/shell.rs

foundry_common/io/
shell.rs
