Skip to main content

plotting/
kitty.rs

1//! Kitty terminal graphics renderer.
2//!
3//! Renders a [`ChartSpec`] to an RGB bitmap via plotters, encodes it as
4//! PNG via the `image` crate, then wraps the PNG in kitty's graphics
5//! Application Programming Command (APC) escape sequence so the result
6//! can be written directly to stdout inside a kitty-compatible terminal.
7//!
8//! Protocol reference: <https://sw.kovidgoyal.net/kitty/graphics-protocol/>.
9//! Key choices:
10//!
11//! - `f=100` — the payload is PNG, not raw RGB. Lets the terminal
12//!   decode via libpng without us having to specify pixel dimensions.
13//! - `a=T` — transmit and display immediately at the cursor.
14//! - `m=1`/`m=0` — chunking. The protocol caps each APC chunk at 4096
15//!   base64 characters; any more gets silently dropped.
16
17use base64::Engine;
18use base64::engine::general_purpose::STANDARD as BASE64;
19use image::{ImageBuffer, Rgb};
20use plotters::prelude::*;
21
22use crate::draw::draw_on;
23use crate::spec::ChartSpec;
24
25/// Dimensions, in pixels, for the rendered chart.
26#[derive(Debug, Clone, Copy)]
27pub struct KittyOpts {
28    pub width_px: u32,
29    pub height_px: u32,
30}
31
32impl Default for KittyOpts {
33    fn default() -> Self {
34        Self {
35            width_px: 640,
36            height_px: 400,
37        }
38    }
39}
40
41/// The kitty graphics protocol caps a single APC chunk's base64 payload
42/// at this many characters.
43const KITTY_CHUNK_MAX: usize = 4096;
44
45/// Render `spec` to a PNG and return the kitty APC escape string that
46/// displays it at the current cursor.
47///
48/// Returns an empty string when drawing fails or when PNG encoding
49/// fails. A blank chart is the least intrusive fallback: the TUI's
50/// chart pane will simply show nothing rather than garbling the
51/// terminal with malformed control bytes.
52#[must_use]
53pub fn render_kitty(spec: &ChartSpec, opts: KittyOpts) -> String {
54    match render_png(spec, opts) {
55        Some(png) => encode_png_apc(&png),
56        None => String::new(),
57    }
58}
59
60/// Render `spec` to a PNG byte vector. Returns `None` on any drawing or
61/// encoding failure; callers decide whether to substitute a text
62/// fallback.
63#[must_use]
64pub fn render_png(spec: &ChartSpec, opts: KittyOpts) -> Option<Vec<u8>> {
65    let width = opts.width_px;
66    let height = opts.height_px;
67    let byte_len = (width as usize)
68        .checked_mul(height as usize)?
69        .checked_mul(3)?;
70    let mut rgb: Vec<u8> = vec![0xff; byte_len];
71
72    let drew = {
73        let backend = BitMapBackend::with_buffer(&mut rgb, (width, height));
74        let root = backend.into_drawing_area();
75        draw_on(&root, spec).is_ok() && root.present().is_ok()
76    };
77
78    if !drew {
79        return None;
80    }
81
82    let buffer: ImageBuffer<Rgb<u8>, _> = ImageBuffer::from_raw(width, height, rgb)?;
83    let mut png = Vec::new();
84    let mut cursor = std::io::Cursor::new(&mut png);
85    buffer.write_to(&mut cursor, image::ImageFormat::Png).ok()?;
86    Some(png)
87}
88
89/// Encode PNG bytes as one or more kitty APC graphics commands. The
90/// chunks share an implicit image ID; the final chunk sets `m=0` to
91/// tell the terminal to display.
92#[must_use]
93pub fn encode_png_apc(png: &[u8]) -> String {
94    let encoded = BASE64.encode(png);
95    let mut out = String::with_capacity(encoded.len() + 64);
96    let chunks: Vec<&str> = encoded
97        .as_bytes()
98        .chunks(KITTY_CHUNK_MAX)
99        .map(|c| std::str::from_utf8(c).unwrap_or(""))
100        .collect();
101
102    if chunks.is_empty() {
103        return String::new();
104    }
105
106    for (i, chunk) in chunks.iter().enumerate() {
107        let first = i == 0;
108        let last = i + 1 == chunks.len();
109        let more = i32::from(!last);
110        if first {
111            out.push_str(&format!("\x1b_Gf=100,a=T,m={more};"));
112        } else {
113            out.push_str(&format!("\x1b_Gm={more};"));
114        }
115        out.push_str(chunk);
116        out.push_str("\x1b\\");
117    }
118    out
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::spec::{ChartKind, Series, SeriesPoint};
125
126    fn sample(kind: ChartKind) -> ChartSpec {
127        ChartSpec {
128            title: "Test".to_string(),
129            kind,
130            x_label: "X".to_string(),
131            y_label: "Y".to_string(),
132            series: vec![Series {
133                label: "A".to_string(),
134                commodity_symbol: "USD".to_string(),
135                points: vec![
136                    SeriesPoint {
137                        x: "Jan".to_string(),
138                        y_num: 100,
139                        y_denom: 1,
140                    },
141                    SeriesPoint {
142                        x: "Feb".to_string(),
143                        y_num: 200,
144                        y_denom: 1,
145                    },
146                ],
147            }],
148            notes: vec![],
149        }
150    }
151
152    #[test]
153    fn render_kitty_wraps_output_in_apc_framing() {
154        let out = render_kitty(
155            &sample(ChartKind::Bar),
156            KittyOpts {
157                width_px: 320,
158                height_px: 200,
159            },
160        );
161        assert!(!out.is_empty(), "kitty output should not be empty");
162        assert!(out.starts_with("\x1b_G"), "missing opening APC marker");
163        assert!(out.ends_with("\x1b\\"), "missing closing APC marker");
164        assert!(out.contains("f=100"), "first chunk must declare PNG format");
165        assert!(out.contains("a=T"), "first chunk must request display");
166    }
167
168    #[test]
169    fn render_png_produces_valid_png_magic_bytes() {
170        let png = render_png(
171            &sample(ChartKind::Line),
172            KittyOpts {
173                width_px: 320,
174                height_px: 200,
175            },
176        )
177        .expect("bitmap render should succeed");
178        assert_eq!(
179            &png[..8],
180            &[0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a],
181            "PNG magic bytes mismatch"
182        );
183    }
184
185    #[test]
186    fn encode_png_apc_round_trips_payload() {
187        let png = b"\x89PNG\r\n\x1a\nsome-fake-payload";
188        let framed = encode_png_apc(png);
189        let payload: String = framed
190            .split("\x1b\\")
191            .filter(|s| !s.is_empty())
192            .map(|chunk| {
193                let semi = chunk.find(';').expect("semicolon separator");
194                &chunk[semi + 1..]
195            })
196            .collect();
197        let decoded = BASE64.decode(payload).expect("decoded base64");
198        assert_eq!(decoded, png);
199    }
200
201    #[test]
202    fn encode_png_apc_chunks_large_payload() {
203        // Produce a PNG whose base64 encoding is > 4096 chars so
204        // chunking logic is exercised.
205        let png: Vec<u8> = (0..4000_u32).flat_map(u32::to_le_bytes).collect();
206        let framed = encode_png_apc(&png);
207        let start_markers = framed.matches("\x1b_G").count();
208        assert!(
209            start_markers >= 2,
210            "expected at least 2 APC chunks, got {start_markers}"
211        );
212        assert!(framed.contains("m=1"), "intermediate chunks must set m=1");
213        assert!(framed.contains("m=0"), "final chunk must set m=0");
214    }
215
216    #[test]
217    fn encode_png_apc_handles_empty_input() {
218        assert!(encode_png_apc(&[]).is_empty());
219    }
220}