Skip to main content

taimux_cli/
ansi.rs

1//! Just enough ANSI to render a captured pane.
2//!
3//! The preview shows a live agent session, which is nothing but colour, so
4//! `capture-pane -e` is what has to be drawn. fzf gets this for free by being
5//! handed the bytes and parsing them itself; a native picker has to do it.
6//!
7//! Deliberately not a terminal emulator, and deliberately not a dependency.
8//! `capture-pane -e` emits SGR (`ESC [ ... m`) and nothing else: the cursor
9//! movement, scrolling and mode changes a real stream carries have already been
10//! applied to the screen it hands back. So this understands SGR, skips any other
11//! escape sequence whole, and never has to be right about anything else.
12
13use ratatui::style::{Color, Modifier, Style};
14use ratatui::text::{Line, Span};
15
16/// One SGR run: the parameters seen so far, applied to a fresh style.
17///
18/// Rebuilt from scratch on each `m` rather than mutated in place, because the
19/// parameters within one sequence are ordered (`0;1;33` is reset, then bold, then
20/// yellow) and an accumulated style has to honour a `0` in the middle of it.
21fn apply(style: Style, params: &[i32]) -> Style {
22    let mut s = style;
23    let mut i = 0;
24    while i < params.len() {
25        let p = params[i];
26        match p {
27            0 => s = Style::default(),
28            1 => s = s.add_modifier(Modifier::BOLD),
29            2 => s = s.add_modifier(Modifier::DIM),
30            3 => s = s.add_modifier(Modifier::ITALIC),
31            4 => s = s.add_modifier(Modifier::UNDERLINED),
32            7 => s = s.add_modifier(Modifier::REVERSED),
33            9 => s = s.add_modifier(Modifier::CROSSED_OUT),
34            21 | 22 => s = s.remove_modifier(Modifier::BOLD | Modifier::DIM),
35            23 => s = s.remove_modifier(Modifier::ITALIC),
36            24 => s = s.remove_modifier(Modifier::UNDERLINED),
37            27 => s = s.remove_modifier(Modifier::REVERSED),
38            29 => s = s.remove_modifier(Modifier::CROSSED_OUT),
39            30..=37 => s = s.fg(basic(p - 30)),
40            39 => s = s.fg(Color::Reset),
41            40..=47 => s = s.bg(basic(p - 40)),
42            49 => s = s.bg(Color::Reset),
43            90..=97 => s = s.fg(bright(p - 90)),
44            100..=107 => s = s.bg(bright(p - 100)),
45            // 38/48 take their colour from the parameters that follow: `5;N` for
46            // one of the 256, `2;r;g;b` for a true colour. Anything else is a
47            // sequence we do not know, and skipping just the introducer would
48            // leave its arguments to be read as colours of their own.
49            38 | 48 => {
50                let fg = p == 38;
51                match params.get(i + 1) {
52                    Some(5) => {
53                        if let Some(&n) = params.get(i + 2) {
54                            let c = Color::Indexed(n.clamp(0, 255) as u8);
55                            s = if fg { s.fg(c) } else { s.bg(c) };
56                        }
57                        i += 2;
58                    }
59                    Some(2) => {
60                        if let (Some(&r), Some(&g), Some(&b)) =
61                            (params.get(i + 2), params.get(i + 3), params.get(i + 4))
62                        {
63                            let c = Color::Rgb(r as u8, g as u8, b as u8);
64                            s = if fg { s.fg(c) } else { s.bg(c) };
65                        }
66                        i += 4;
67                    }
68                    _ => break,
69                }
70            }
71            _ => {}
72        }
73        i += 1;
74    }
75    s
76}
77
78fn basic(n: i32) -> Color {
79    match n {
80        0 => Color::Black,
81        1 => Color::Red,
82        2 => Color::Green,
83        3 => Color::Yellow,
84        4 => Color::Blue,
85        5 => Color::Magenta,
86        6 => Color::Cyan,
87        _ => Color::Gray,
88    }
89}
90
91fn bright(n: i32) -> Color {
92    match n {
93        0 => Color::DarkGray,
94        1 => Color::LightRed,
95        2 => Color::LightGreen,
96        3 => Color::LightYellow,
97        4 => Color::LightBlue,
98        5 => Color::LightMagenta,
99        6 => Color::LightCyan,
100        _ => Color::White,
101    }
102}
103
104/// Text into styled lines, carrying the style across line breaks the way a
105/// terminal does: a capture can open a colour on one row and close it on the next.
106pub fn to_lines(text: &str) -> Vec<Line<'static>> {
107    let mut lines = Vec::new();
108    let mut style = Style::default();
109    for raw in text.lines() {
110        let mut spans: Vec<Span<'static>> = Vec::new();
111        let mut buf = String::new();
112        let mut it = raw.char_indices().peekable();
113        while let Some((_, c)) = it.next() {
114            if c != '\x1b' {
115                buf.push(c);
116                continue;
117            }
118            // Everything up to the escape keeps the style in force.
119            if !buf.is_empty() {
120                spans.push(Span::styled(std::mem::take(&mut buf), style));
121            }
122            match it.peek().map(|&(_, c)| c) {
123                Some('[') => {
124                    it.next();
125                    // A CSI is parameter bytes (0x30-0x3F, so digits and `;` but
126                    // also the private markers `?<=>`), then intermediate bytes
127                    // (0x20-0x2F), then one final byte (0x40-0x7E). Stopping at
128                    // the first non-digit instead treats the `?` of `ESC[?25l` as
129                    // the end of the sequence and prints "25l" on the screen.
130                    let mut params: Vec<i32> = Vec::new();
131                    let mut num = String::new();
132                    let mut private = false;
133                    let mut kind = None;
134                    for (_, c) in it.by_ref() {
135                        match c {
136                            '0'..='9' => num.push(c),
137                            ';' | ':' => {
138                                params.push(num.parse().unwrap_or(0));
139                                num.clear();
140                            }
141                            '?' | '<' | '=' | '>' => private = true,
142                            ' '..='/' => {} // intermediate
143                            _ => {
144                                kind = Some(c);
145                                break;
146                            }
147                        }
148                    }
149                    if !num.is_empty() || params.is_empty() {
150                        params.push(num.parse().unwrap_or(0));
151                    }
152                    // Only SGR changes anything. Any other final byte is a
153                    // sequence this does not speak, and it is dropped rather than
154                    // printed, which is the whole job.
155                    if kind == Some('m') && !private {
156                        style = apply(style, &params);
157                    }
158                }
159                // OSC and the rest: swallow to the terminator so the payload does
160                // not land on screen as text.
161                Some(']') => {
162                    for (_, c) in it.by_ref() {
163                        if c == '\x07' || c == '\x1b' {
164                            break;
165                        }
166                    }
167                }
168                _ => {
169                    it.next();
170                }
171            }
172        }
173        if !buf.is_empty() {
174            spans.push(Span::styled(buf, style));
175        }
176        lines.push(Line::from(spans));
177    }
178    lines
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn texts(l: &Line) -> Vec<String> {
186        l.spans.iter().map(|s| s.content.to_string()).collect()
187    }
188
189    #[test]
190    fn plain_text_is_one_span() {
191        let l = to_lines("hello");
192        assert_eq!(texts(&l[0]), vec!["hello"]);
193        assert_eq!(l[0].spans[0].style, Style::default());
194    }
195
196    #[test]
197    fn a_colour_opens_a_span_and_a_reset_closes_it() {
198        let l = to_lines("a\x1b[31mred\x1b[0mb");
199        assert_eq!(texts(&l[0]), vec!["a", "red", "b"]);
200        assert_eq!(l[0].spans[1].style.fg, Some(Color::Red));
201        assert_eq!(l[0].spans[2].style, Style::default());
202    }
203
204    /// The parameters within one sequence are ordered, so a reset in the middle
205    /// of one has to clear what came before it.
206    #[test]
207    fn a_reset_inside_a_sequence_clears_what_preceded_it() {
208        let l = to_lines("\x1b[1;33mx\x1b[0;36my");
209        assert!(l[0].spans[0].style.add_modifier.contains(Modifier::BOLD));
210        assert_eq!(l[0].spans[1].style.fg, Some(Color::Cyan));
211        assert!(!l[0].spans[1].style.add_modifier.contains(Modifier::BOLD));
212    }
213
214    #[test]
215    fn indexed_and_true_colour_both_land() {
216        let l = to_lines("\x1b[38;5;208mo\x1b[38;2;18;52;86mt");
217        assert_eq!(l[0].spans[0].style.fg, Some(Color::Indexed(208)));
218        assert_eq!(l[0].spans[1].style.fg, Some(Color::Rgb(18, 52, 86)));
219    }
220
221    /// The arguments of a 256-colour sequence must not be read as colours in
222    /// their own right: `38;5;1` is orange-ish 1, not red then something.
223    #[test]
224    fn the_arguments_of_an_extended_colour_are_consumed() {
225        let l = to_lines("\x1b[38;5;1;1mx");
226        assert_eq!(l[0].spans[0].style.fg, Some(Color::Indexed(1)));
227        assert!(l[0].spans[0].style.add_modifier.contains(Modifier::BOLD));
228    }
229
230    /// Style carries across a line break, the way a terminal does: a capture can
231    /// open a colour on one row and close it on the next.
232    #[test]
233    fn style_carries_from_one_line_to_the_next() {
234        let l = to_lines("\x1b[32mgreen\nstill green\x1b[0m\nplain");
235        assert_eq!(l[1].spans[0].style.fg, Some(Color::Green));
236        assert_eq!(l[2].spans[0].style, Style::default());
237    }
238
239    /// Anything that is not SGR is dropped rather than printed. A capture should
240    /// not contain cursor movement, but a pane holding a raw log might.
241    #[test]
242    fn a_non_sgr_sequence_is_swallowed_whole() {
243        let l = to_lines("a\x1b[2Jb\x1b[?25lc");
244        assert_eq!(texts(&l[0]).concat(), "abc");
245    }
246
247    #[test]
248    fn an_osc_title_does_not_land_on_screen() {
249        let l = to_lines("a\x1b]0;a window title\x07b");
250        assert_eq!(texts(&l[0]).concat(), "ab");
251    }
252
253    #[test]
254    fn a_bare_escape_at_the_end_does_not_panic() {
255        assert_eq!(to_lines("a\x1b").len(), 1);
256        assert_eq!(to_lines("\x1b[").len(), 1);
257        assert_eq!(to_lines("\x1b[38;5").len(), 1);
258    }
259
260    /// Blank lines are kept: a captured screen is padded to the pane height, and
261    /// dropping the empties would close gaps the session actually has.
262    #[test]
263    fn blank_lines_survive() {
264        assert_eq!(to_lines("a\n\nb").len(), 3);
265    }
266}