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

foundry_common/io/
shell.rs
