anvil_zksync_common/
shell.rs

1//! Utility functions for writing to stdout and stderr.
2//!
3//! Simplified adaptation from
4//! https://github.com/foundry-rs/foundry/blob/master/crates/common/src/io/macros.rs.
5
6use anstream::{AutoStream, ColorChoice as AnstreamColorChoice};
7use anstyle::{AnsiColor, Effects, Reset, Style};
8use std::fmt::Arguments;
9use std::io::{self, IsTerminal, Write};
10use std::sync::{Mutex, OnceLock};
11
12/// Default styles for WARNING and ERROR messages.
13pub const ERROR: Style = AnsiColor::Red.on_default().effects(Effects::BOLD);
14pub const WARN: Style = AnsiColor::Yellow.on_default().effects(Effects::BOLD);
15
16/// Choices for whether to use colored output.
17#[derive(Debug, Clone, Copy, PartialEq)]
18pub enum ColorChoice {
19    Auto,
20    Always,
21    Never,
22}
23
24impl ColorChoice {
25    #[inline]
26    fn to_anstream(self) -> AnstreamColorChoice {
27        match self {
28            ColorChoice::Always => AnstreamColorChoice::Always,
29            ColorChoice::Never => AnstreamColorChoice::Never,
30            ColorChoice::Auto => AnstreamColorChoice::Auto,
31        }
32    }
33}
34
35/// The output mode: either normal output or completely quiet.
36#[derive(Debug, Clone, Copy, PartialEq)]
37pub enum OutputMode {
38    Normal,
39    Quiet,
40}
41
42/// A shell abstraction that tracks verbosity level, an output mode,
43/// and a color choice.
44/// It uses anstream’s AutoStream for stdout and stderr.
45#[derive(Debug)]
46pub struct Shell {
47    /// Verbosity level.
48    pub verbosity: u8,
49    /// Whether to output anything at all.
50    pub output_mode: OutputMode,
51    /// Whether to use colors.
52    pub color_choice: ColorChoice,
53    /// Adaptive stdout.
54    stdout: AutoStream<std::io::Stdout>,
55    /// Adaptive stderr.
56    stderr: AutoStream<std::io::Stderr>,
57}
58
59impl Shell {
60    /// Creates a new shell with default settings:
61    /// - verbosity 0,
62    /// - Normal output mode,
63    /// - Auto color detection.
64    pub fn new() -> Self {
65        let color = ColorChoice::Auto;
66        Self {
67            verbosity: 0,
68            output_mode: OutputMode::Normal,
69            color_choice: color,
70            stdout: AutoStream::new(std::io::stdout(), color.to_anstream()),
71            stderr: AutoStream::new(std::io::stderr(), color.to_anstream()),
72        }
73    }
74
75    /// Print a string to stdout.
76    pub fn print_out(&mut self, args: Arguments) -> io::Result<()> {
77        if self.output_mode == OutputMode::Quiet {
78            return Ok(());
79        }
80
81        self.stdout.write_fmt(args)?;
82        self.stdout.flush()
83    }
84
85    /// Print a line (with a newline) to stdout.
86    pub fn println_out(&mut self, args: Arguments) -> io::Result<()> {
87        if self.output_mode == OutputMode::Quiet {
88            return Ok(());
89        }
90
91        self.stdout.write_fmt(args)?;
92        writeln!(self.stdout)?;
93        self.stdout.flush()
94    }
95
96    /// Print a string to stderr.
97    pub fn print_err(&mut self, args: Arguments) -> io::Result<()> {
98        if self.output_mode == OutputMode::Quiet {
99            return Ok(());
100        }
101        self.stderr.write_fmt(args)?;
102        self.stderr.flush()
103    }
104
105    /// Print a line (with a newline) to stderr.
106    pub fn println_err(&mut self, args: Arguments) -> io::Result<()> {
107        if self.output_mode == OutputMode::Quiet {
108            return Ok(());
109        }
110        self.stderr.write_fmt(args)?;
111        writeln!(self.stderr)?;
112        self.stderr.flush()
113    }
114
115    /// Print a warning message.
116    ///
117    /// If colors are enabled, the “Warning:” prefix is printed in yellow.
118    pub fn warn(&mut self, args: Arguments) -> io::Result<()> {
119        if self.should_color() {
120            write!(self.stderr, "{}Warning:{} ", WARN, Reset)?;
121        } else {
122            write!(self.stderr, "Warning: ")?;
123        }
124        self.stderr.write_fmt(args)?;
125        writeln!(self.stderr)?;
126        self.stderr.flush()
127    }
128
129    /// Print an error message.
130    ///
131    /// If colors are enabled, the “Error:” prefix is printed in red.
132    pub fn error(&mut self, args: Arguments) -> io::Result<()> {
133        if self.should_color() {
134            write!(self.stderr, "{}Error:{} ", ERROR, Reset)?;
135        } else {
136            write!(self.stderr, "Error: ")?;
137        }
138        self.stderr.write_fmt(args)?;
139        writeln!(self.stderr)?;
140        self.stderr.flush()
141    }
142
143    fn should_color(&self) -> bool {
144        match self.color_choice {
145            ColorChoice::Always => true,
146            ColorChoice::Never => false,
147            ColorChoice::Auto => std::io::stdout().is_terminal(),
148        }
149    }
150}
151
152impl Default for Shell {
153    fn default() -> Self {
154        Self::new()
155    }
156}
157
158/// The global shell instance.
159///
160/// This uses a [`OnceLock`](https://doc.rust-lang.org/std/sync/struct.OnceLock.html)
161/// to initialize the shell once and then a [`Mutex`](std::sync::Mutex) to allow
162/// mutable access to it from anywhere.
163static GLOBAL_SHELL: OnceLock<Mutex<Shell>> = OnceLock::new();
164
165/// Get a lock to the global shell.
166///
167/// This will initialize the shell with default values if it has not been set yet.
168pub fn get_shell() -> std::sync::MutexGuard<'static, Shell> {
169    GLOBAL_SHELL
170        .get_or_init(|| Mutex::new(Shell::new()))
171        .lock()
172        .expect("global shell mutex is poisoned")
173}
174
175/// (Optional) Set the global shell with a custom configuration.
176///
177/// Note that this will fail if the shell has already been set.
178pub fn set_shell(shell: Shell) {
179    let _ = GLOBAL_SHELL.set(Mutex::new(shell));
180}
181
182/// Macro to print a formatted message to stdout.
183///
184/// Usage:
185/// ```
186/// use anvil_zksync_common::sh_print;
187///
188/// sh_print!("Hello, {}!", "world");
189/// ```
190#[macro_export]
191macro_rules! sh_print {
192    ($($arg:tt)*) => {{
193        $crate::shell::get_shell().print_out(format_args!($($arg)*))
194            .unwrap_or_else(|e| eprintln!("Error writing output: {}", e));
195    }};
196}
197
198/// Macro to print a formatted message (with a newline) to stdout.
199#[macro_export]
200macro_rules! sh_println {
201    ($($arg:tt)*) => {{
202        $crate::shell::get_shell().println_out(format_args!($($arg)*))
203            .unwrap_or_else(|e| eprintln!("Error writing output: {}", e));
204    }};
205}
206
207/// Macro to print a formatted message to stderr.
208#[macro_export]
209macro_rules! sh_eprint {
210    ($($arg:tt)*) => {{
211        $crate::shell::get_shell().print_err(format_args!($($arg)*))
212            .unwrap_or_else(|e| eprintln!("Error writing stderr: {}", e));
213    }};
214}
215
216/// Macro to print a formatted message (with newline) to stderr.
217#[macro_export]
218macro_rules! sh_eprintln {
219    ($($arg:tt)*) => {{
220        $crate::shell::get_shell().println_err(format_args!($($arg)*))
221            .unwrap_or_else(|e| eprintln!("Error writing stderr: {}", e));
222    }};
223}
224
225/// Macro to print a warning message.
226///
227/// Usage:
228/// ```
229/// use anvil_zksync_common::sh_warn;
230///
231/// sh_warn!("This is a warning: {}", "be careful!");
232/// ```
233#[macro_export]
234macro_rules! sh_warn {
235    ($($arg:tt)*) => {{
236        $crate::shell::get_shell().warn(format_args!($($arg)*))
237            .unwrap_or_else(|e| eprintln!("Error writing warning: {}", e));
238    }};
239}
240
241/// Macro to print an error message.
242///
243/// Usage:
244/// ```
245/// use anvil_zksync_common::sh_err;
246///
247/// sh_err!("Something went wrong: {}", "details");
248/// ```
249#[macro_export]
250macro_rules! sh_err {
251    ($($arg:tt)*) => {{
252        $crate::shell::get_shell().error(format_args!($($arg)*))
253            .unwrap_or_else(|e| eprintln!("Error writing error: {}", e));
254    }};
255}
256
257#[cfg(test)]
258mod tests {
259    #[test]
260    fn test_shell_macros() {
261        sh_print!("Hello, ");
262        sh_println!("world!");
263        sh_eprint!("Error: ");
264        sh_eprintln!("Something went wrong!");
265        sh_warn!("This is a warning");
266        sh_err!("This is an error");
267    }
268}