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

foundry_common/io/
shell.rs
