Skip to main content

nms/
graphics.rs

1//! Inline kitty graphics support for the `nms` REPL / one-shot
2//! evaluator.
3//!
4//! When the user evaluates a form whose value is a `Value::Bytes`
5//! that sniffs as an image AND stdout is connected to a kitty-graphics
6//! capable terminal, the value-printer emits a kitty APC sequence
7//! containing the image. The image renders inline at the cursor;
8//! subsequent text lines below sit beneath it.
9//!
10//! Capability detection mirrors `tui::chart::supports_kitty` so an
11//! `nms` invocation and a TUI session inside the same terminal agree
12//! on what to emit. The `--no-graphics` CLI flag short-circuits the
13//! detection for CI captures and for users whose terminal the
14//! heuristics misidentify.
15
16use base64::Engine;
17use base64::engine::general_purpose::STANDARD;
18
19/// Image format the kitty graphics protocol can transmit directly.
20/// `f=100` corresponds to PNG; `f=24`/`f=32` would be raw RGB/RGBA.
21/// We never auto-convert SVG → PNG inline (that lives in the
22/// `plotting` crate's `:format png` keyword path); SVG bytes fall
23/// through to the textual `#u8(...)` printer.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum ImageFormat {
26    Png,
27    Jpeg,
28}
29
30impl ImageFormat {
31    fn kitty_format_code(self) -> &'static str {
32        match self {
33            // Kitty's f=100 covers both PNG and JPEG decode; the
34            // protocol picks the right decoder via libpng/libjpeg.
35            ImageFormat::Png | ImageFormat::Jpeg => "100",
36        }
37    }
38}
39
40/// Detect the image format by magic bytes. Returns `None` for any
41/// payload that isn't a known inline-renderable raster.
42#[must_use]
43pub fn sniff(bytes: &[u8]) -> Option<ImageFormat> {
44    if bytes.len() >= 8 && bytes[0..8] == [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] {
45        return Some(ImageFormat::Png);
46    }
47    if bytes.len() >= 3 && bytes[0..3] == [0xFF, 0xD8, 0xFF] {
48        return Some(ImageFormat::Jpeg);
49    }
50    None
51}
52
53/// Mirror of `tui::chart::supports_kitty` for the standalone `nms`
54/// process. Returns true when the surrounding terminal advertises
55/// kitty-graphics support via env vars set by the terminal emulator.
56#[must_use]
57pub fn supports_kitty<F>(env_lookup: F) -> bool
58where
59    F: Fn(&str) -> Option<String>,
60{
61    if env_lookup("KITTY_WINDOW_ID").is_some() {
62        return true;
63    }
64    let term = env_lookup("TERM").unwrap_or_default();
65    if term.contains("kitty") {
66        return true;
67    }
68    let term_program = env_lookup("TERM_PROGRAM").unwrap_or_default();
69    matches!(term_program.as_str(), "WezTerm" | "ghostty")
70}
71
72/// Wrap `bytes` of a known image format into a single kitty graphics
73/// APC sequence. The sequence is appended to `out`, which keeps
74/// callers in control of stdout buffering — `nms` prints it through
75/// the normal value-formatter, and tests can capture the bytes for
76/// comparison.
77///
78/// For payloads above 4096 bytes the protocol recommends chunking with
79/// the `m=1`/`m=0` continuation markers. We chunk at 4096 base64
80/// characters (≈ 3072 raw bytes) so even multi-MB charts transmit
81/// without exceeding terminal buffer limits.
82pub fn encode_kitty(format: ImageFormat, bytes: &[u8], out: &mut String) {
83    let encoded = STANDARD.encode(bytes);
84    const CHUNK: usize = 4096;
85    let f = format.kitty_format_code();
86
87    if encoded.len() <= CHUNK {
88        out.push_str("\x1b_Ga=T,f=");
89        out.push_str(f);
90        out.push(';');
91        out.push_str(&encoded);
92        out.push_str("\x1b\\");
93        return;
94    }
95
96    let mut i = 0;
97    let mut first = true;
98    while i < encoded.len() {
99        let end = (i + CHUNK).min(encoded.len());
100        let last = end == encoded.len();
101        out.push_str("\x1b_G");
102        if first {
103            out.push_str("a=T,f=");
104            out.push_str(f);
105            out.push(',');
106            first = false;
107        }
108        out.push_str(if last { "m=0;" } else { "m=1;" });
109        out.push_str(&encoded[i..end]);
110        out.push_str("\x1b\\");
111        i = end;
112    }
113}
114
115/// Convenience entry that returns the kitty APC payload as a String
116/// when the bytes are a known image and the terminal supports kitty
117/// graphics, otherwise `None`. The caller falls back to the textual
118/// `#u8(...)` printer when this returns `None`.
119#[must_use]
120pub fn try_render_inline<F>(bytes: &[u8], env_lookup: F) -> Option<String>
121where
122    F: Fn(&str) -> Option<String>,
123{
124    let format = sniff(bytes)?;
125    if !supports_kitty(env_lookup) {
126        return None;
127    }
128    let mut out = String::with_capacity(bytes.len() * 4 / 3 + 32);
129    encode_kitty(format, bytes, &mut out);
130    Some(out)
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use std::collections::HashMap;
137
138    fn env_from(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
139        let map: HashMap<String, String> = pairs
140            .iter()
141            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
142            .collect();
143        move |k: &str| map.get(k).cloned()
144    }
145
146    #[test]
147    fn sniff_recognises_png_magic() {
148        let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0, 0];
149        assert_eq!(sniff(&png), Some(ImageFormat::Png));
150    }
151
152    #[test]
153    fn sniff_recognises_jpeg_magic() {
154        let jpeg = [0xFF, 0xD8, 0xFF, 0xE0, 0, 0];
155        assert_eq!(sniff(&jpeg), Some(ImageFormat::Jpeg));
156    }
157
158    #[test]
159    fn sniff_returns_none_for_random_bytes() {
160        assert_eq!(sniff(&[1, 2, 3, 4]), None);
161        assert_eq!(sniff(b"<svg ..."), None);
162        assert_eq!(sniff(&[]), None);
163    }
164
165    #[test]
166    fn supports_kitty_window_id() {
167        let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
168        assert!(supports_kitty(env));
169    }
170
171    #[test]
172    fn supports_kitty_via_term() {
173        let env = env_from(&[("TERM", "xterm-kitty")]);
174        assert!(supports_kitty(env));
175    }
176
177    #[test]
178    fn supports_kitty_via_term_program_wezterm() {
179        let env = env_from(&[("TERM_PROGRAM", "WezTerm")]);
180        assert!(supports_kitty(env));
181    }
182
183    #[test]
184    fn supports_kitty_via_term_program_ghostty() {
185        let env = env_from(&[("TERM_PROGRAM", "ghostty")]);
186        assert!(supports_kitty(env));
187    }
188
189    #[test]
190    fn supports_kitty_plain_xterm_is_false() {
191        let env = env_from(&[("TERM", "xterm-256color")]);
192        assert!(!supports_kitty(env));
193    }
194
195    #[test]
196    fn supports_kitty_empty_env_is_false() {
197        let env = env_from(&[]);
198        assert!(!supports_kitty(env));
199    }
200
201    #[test]
202    fn encode_kitty_short_payload_uses_single_chunk() {
203        let mut out = String::new();
204        encode_kitty(ImageFormat::Png, &[0x89, 0x50, 0x4E, 0x47], &mut out);
205        // Single-chunk sequence: ESC _ G a=T,f=100;<b64>ESC\
206        assert!(out.starts_with("\x1b_Ga=T,f=100;"));
207        assert!(out.ends_with("\x1b\\"));
208        // Exactly one APC start sequence.
209        assert_eq!(out.matches("\x1b_G").count(), 1);
210    }
211
212    #[test]
213    fn encode_kitty_long_payload_chunks() {
214        let bytes = vec![0xABu8; 10_000];
215        let mut out = String::new();
216        encode_kitty(ImageFormat::Png, &bytes, &mut out);
217        // Expect multiple APC chunks; the first carries the format
218        // header, the last has m=0 (terminator).
219        let starts = out.matches("\x1b_G").count();
220        assert!(starts > 1, "got {starts} chunks");
221        assert!(
222            out.ends_with(
223                "m=0;"
224                    .chars()
225                    .chain("".chars())
226                    .collect::<String>()
227                    .as_str()
228            ) || out.contains("m=0;")
229        );
230    }
231
232    #[test]
233    fn try_render_inline_returns_none_for_non_image() {
234        let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
235        assert!(try_render_inline(b"not an image", env).is_none());
236    }
237
238    #[test]
239    fn try_render_inline_returns_none_when_terminal_lacks_kitty() {
240        let env = env_from(&[("TERM", "xterm-256color")]);
241        let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
242        assert!(try_render_inline(&png, env).is_none());
243    }
244
245    #[test]
246    fn try_render_inline_returns_some_when_png_and_kitty() {
247        let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
248        let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
249        let rendered = try_render_inline(&png, env).expect("should encode");
250        assert!(rendered.contains("\x1b_G"));
251    }
252}