command_stream/zx/
log.rs

1//! Verbose logging (zx `log`) and the command highlighter (zx `formatCmd`).
2//!
3//! Entries are written only when they are verbose. Colors follow zx: the
4//! command name is bright green, operators red, quoted strings bright yellow
5//! and shell keywords bright cyan. [`format_cmd`] always colors (like zx with
6//! a forced color level); [`Logger::stderr`] enables colors only when stderr
7//! is a terminal and `NO_COLOR` is unset.
8
9use std::collections::HashMap;
10use std::io::{IsTerminal, Write};
11use std::time::Duration;
12
13/// A pair of ANSI open/close sequences.
14#[derive(Debug, Clone, Copy)]
15struct Style {
16    open: &'static str,
17    close: &'static str,
18}
19
20const RED: Style = Style {
21    open: "\x1b[31m",
22    close: "\x1b[39m",
23};
24const GREEN_BRIGHT: Style = Style {
25    open: "\x1b[92m",
26    close: "\x1b[39m",
27};
28const YELLOW_BRIGHT: Style = Style {
29    open: "\x1b[93m",
30    close: "\x1b[39m",
31};
32const CYAN_BRIGHT: Style = Style {
33    open: "\x1b[96m",
34    close: "\x1b[39m",
35};
36const RESET: Style = Style {
37    open: "\x1b[0m",
38    close: "\x1b[0m",
39};
40const FAIL_BADGE: Style = Style {
41    open: "\x1b[41m\x1b[37m",
42    close: "\x1b[39m\x1b[49m",
43};
44
45/// Wrap `text` in `style`, closing and reopening the style around line
46/// breaks so that a prefix inserted after `\n` stays uncolored.
47fn paint(style: Style, text: &str, colors: bool) -> String {
48    if !colors || text.is_empty() {
49        return text.to_string();
50    }
51    let mut out = String::with_capacity(text.len() + 16);
52    out.push_str(style.open);
53    let mut rest = text;
54    while let Some(i) = rest.find('\n') {
55        let cut = if rest[..i].ends_with('\r') { i - 1 } else { i };
56        out.push_str(&rest[..cut]);
57        out.push_str(style.close);
58        out.push_str(&rest[cut..=i]);
59        out.push_str(style.open);
60        rest = &rest[i + 1..];
61    }
62    out.push_str(rest);
63    out.push_str(style.close);
64    out
65}
66
67const SYNTAX: &str = "()[]{}<>;:+|&=";
68const CMD_BREAK: &str = "|&;><";
69const RESERVED_WORDS: [&str; 15] = [
70    "if", "then", "else", "elif", "fi", "case", "esac", "for", "select", "while", "until", "do",
71    "done", "in", "EOF",
72];
73
74#[derive(Clone, Copy, PartialEq, Eq)]
75enum Mode {
76    Plain,
77    Quote,
78    Dollar,
79    Syntax,
80}
81
82/// Word position inside a simple command: `First` is the command name.
83#[derive(Clone, Copy, PartialEq, Eq)]
84enum Position {
85    Start,
86    First,
87    Assignment,
88    Rest,
89}
90
91struct Highlighter {
92    colors: bool,
93    out: String,
94    word: String,
95    mode: Mode,
96    quote: Option<char>,
97    pos: Position,
98}
99
100impl Highlighter {
101    fn advance(&mut self) {
102        self.pos = match self.pos {
103            Position::Start => Position::First,
104            Position::Assignment => Position::Start,
105            _ => Position::Rest,
106        };
107    }
108
109    fn style_word(&mut self, trimmed: &str) -> String {
110        self.advance();
111        match self.mode {
112            Mode::Syntax => {
113                if CMD_BREAK.contains(trimmed) {
114                    self.pos = Position::Start;
115                }
116                paint(RED, &self.word, self.colors)
117            }
118            Mode::Quote | Mode::Dollar => paint(YELLOW_BRIGHT, &self.word, self.colors),
119            Mode::Plain if RESERVED_WORDS.contains(&trimmed) => {
120                paint(CYAN_BRIGHT, &self.word, self.colors)
121            }
122            Mode::Plain if self.pos == Position::First => {
123                self.pos = Position::Rest;
124                paint(GREEN_BRIGHT, &self.word, self.colors)
125            }
126            Mode::Plain => self.word.clone(),
127        }
128    }
129
130    fn flush(&mut self) {
131        let trimmed = self.word.trim().to_string();
132        let piece = if trimmed.is_empty() {
133            self.word.clone()
134        } else {
135            self.style_word(&trimmed)
136        };
137        self.out.push_str(&piece);
138        self.word.clear();
139        self.mode = Mode::Plain;
140    }
141
142    fn emit_single(&mut self, mode: Mode, c: char) {
143        self.flush();
144        self.mode = mode;
145        self.word.push(c);
146        self.flush();
147    }
148
149    fn push(&mut self, c: char) {
150        if let Some(q) = self.quote {
151            self.word.push(c);
152            if c == q {
153                self.flush();
154                self.quote = None;
155            }
156        } else if c == '$' {
157            self.emit_single(Mode::Dollar, c);
158        } else if c == '\'' || c == '"' {
159            self.flush();
160            self.mode = Mode::Quote;
161            self.quote = Some(c);
162            self.word.push(c);
163        } else if c.is_whitespace() {
164            self.flush();
165            self.word.push(c);
166        } else if SYNTAX.contains(c) {
167            // `FOO=bar cmd`: an assignment before the command keeps the
168            // command-name color for the word after it.
169            let assignment = c == '=' && self.pos == Position::Start;
170            if assignment {
171                // The pending word (the variable name) must not be painted
172                // as the command name.
173                self.pos = Position::First;
174            }
175            self.emit_single(Mode::Syntax, c);
176            if assignment {
177                self.pos = Position::Assignment;
178            }
179        } else {
180            self.word.push(c);
181        }
182    }
183}
184
185fn highlight(cmd: &str, colors: bool) -> String {
186    let mut h = Highlighter {
187        colors,
188        out: String::new(),
189        word: String::new(),
190        mode: Mode::Plain,
191        quote: None,
192        pos: Position::Start,
193    };
194    for c in cmd.chars() {
195        h.push(c);
196    }
197    h.flush();
198    let continuation = paint(RESET, "\n> ", colors);
199    format!("$ {}\n", h.out.replace('\n', &continuation))
200}
201
202/// Render a command as zx prints it in verbose mode (`$ cmd`, continuation
203/// lines prefixed with `> `), with ANSI colors.
204pub fn format_cmd(cmd: &str) -> String {
205    highlight(cmd, true)
206}
207
208/// Like [`format_cmd`] without colors.
209pub fn format_cmd_plain(cmd: &str) -> String {
210    highlight(cmd, false)
211}
212
213/// A log record (zx `LogEntry`).
214#[derive(Debug, Clone, PartialEq)]
215pub enum LogEntry {
216    /// A command about to run.
217    Cmd {
218        /// The command line.
219        cmd: String,
220    },
221    /// A chunk of process stdout.
222    Stdout {
223        /// Raw bytes.
224        data: Vec<u8>,
225    },
226    /// A chunk of process stderr.
227    Stderr {
228        /// Raw bytes.
229        data: Vec<u8>,
230    },
231    /// A directory change.
232    Cd {
233        /// The new directory.
234        dir: String,
235    },
236    /// An HTTP request.
237    Fetch {
238        /// Target URL.
239        url: String,
240        /// Rendered request options, if any.
241        init: Option<String>,
242    },
243    /// Arbitrary text.
244    Custom {
245        /// Text written as-is.
246        data: String,
247    },
248    /// A failed retry attempt.
249    Retry {
250        /// 1-based attempt number.
251        attempt: usize,
252        /// Total attempts (`None` for unbounded).
253        total: Option<usize>,
254        /// Delay before the next attempt.
255        delay: Duration,
256    },
257    /// A process finished.
258    End {
259        /// Exit code, if any.
260        exit_code: Option<i32>,
261        /// Terminating signal, if any.
262        signal: Option<String>,
263        /// Run time.
264        duration: Duration,
265    },
266    /// A signal was sent.
267    Kill {
268        /// Target pid.
269        pid: u32,
270        /// Signal name, if any.
271        signal: Option<String>,
272    },
273}
274
275impl LogEntry {
276    /// The zx `kind` of the entry (`cmd`, `stdout`, `retry`, ...).
277    pub fn kind(&self) -> &'static str {
278        match self {
279            LogEntry::Cmd { .. } => "cmd",
280            LogEntry::Stdout { .. } => "stdout",
281            LogEntry::Stderr { .. } => "stderr",
282            LogEntry::Cd { .. } => "cd",
283            LogEntry::Fetch { .. } => "fetch",
284            LogEntry::Custom { .. } => "custom",
285            LogEntry::Retry { .. } => "retry",
286            LogEntry::End { .. } => "end",
287            LogEntry::Kill { .. } => "kill",
288        }
289    }
290}
291
292/// The default rendering of `entry` (empty for `end` and `kill`).
293pub fn format_entry(entry: &LogEntry, colors: bool) -> Vec<u8> {
294    let text = match entry {
295        LogEntry::Cmd { cmd } => highlight(cmd, colors),
296        LogEntry::Stdout { data } | LogEntry::Stderr { data } => return data.clone(),
297        LogEntry::Custom { data } => data.clone(),
298        LogEntry::Cd { dir } => format!("$ {} {dir}\n", paint(GREEN_BRIGHT, "cd", colors)),
299        LogEntry::Fetch { url, init } => {
300            let init = init.as_ref().map(|i| format!(" {i}")).unwrap_or_default();
301            format!("$ {} {url}{init}\n", paint(GREEN_BRIGHT, "fetch", colors))
302        }
303        LogEntry::Retry {
304            attempt,
305            total,
306            delay,
307        } => {
308            let total = total.map(|t| format!("/{t}")).unwrap_or_default();
309            let delay = if delay.is_zero() {
310                String::new()
311            } else {
312                format!("; next in {}ms", delay.as_millis())
313            };
314            let badge = paint(FAIL_BADGE, " FAIL ", colors);
315            format!("{badge} Attempt: {attempt}{total}{delay}\n")
316        }
317        LogEntry::End { .. } | LogEntry::Kill { .. } => String::new(),
318    };
319    text.into_bytes()
320}
321
322/// A custom renderer registered with [`Logger::formatter`].
323pub type Formatter = Box<dyn Fn(&LogEntry) -> String + Send + Sync>;
324
325/// Writes verbose [`LogEntry`] records to an output (zx `log`).
326pub struct Logger<W: Write> {
327    output: W,
328    colors: bool,
329    formatters: HashMap<&'static str, Formatter>,
330}
331
332impl Logger<std::io::Stderr> {
333    /// A logger writing to stderr, colored when stderr is a terminal and
334    /// `NO_COLOR` is unset.
335    pub fn stderr() -> Self {
336        let colors = std::io::stderr().is_terminal() && std::env::var_os("NO_COLOR").is_none();
337        Logger::new(std::io::stderr()).colors(colors)
338    }
339}
340
341impl<W: Write> Logger<W> {
342    /// A logger writing colored output to `output`.
343    pub fn new(output: W) -> Self {
344        Self {
345            output,
346            colors: true,
347            formatters: HashMap::new(),
348        }
349    }
350
351    /// Enable or disable ANSI colors.
352    pub fn colors(mut self, colors: bool) -> Self {
353        self.colors = colors;
354        self
355    }
356
357    /// Override the rendering of one entry kind (zx `log.formatters`).
358    pub fn formatter(
359        mut self,
360        kind: &'static str,
361        f: impl Fn(&LogEntry) -> String + Send + Sync + 'static,
362    ) -> Self {
363        self.formatters.insert(kind, Box::new(f));
364        self
365    }
366
367    /// The underlying output.
368    pub fn output(&self) -> &W {
369        &self.output
370    }
371
372    /// Consume the logger, returning its output.
373    pub fn into_output(self) -> W {
374        self.output
375    }
376
377    /// Write `entry` if `verbose` is set.
378    pub fn log(&mut self, entry: &LogEntry, verbose: bool) -> std::io::Result<()> {
379        if !verbose {
380            return Ok(());
381        }
382        let bytes = match self.formatters.get(entry.kind()) {
383            Some(f) => f(entry).into_bytes(),
384            None => format_entry(entry, self.colors),
385        };
386        if bytes.is_empty() {
387            return Ok(());
388        }
389        self.output.write_all(&bytes)?;
390        self.output.flush()
391    }
392}
393
394/// Log `entry` to stderr when `verbose` is set (errors are ignored).
395pub fn log(entry: &LogEntry, verbose: bool) {
396    if verbose {
397        let _ = Logger::stderr().log(entry, true);
398    }
399}
command_stream/zx/log.rs

command_stream/zx/
log.rs
