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

foundry_common/io/
shell.rs
