foundry_common/io/
shell.rs

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

foundry_common/io/
shell.rs
