onefetch_image/
kitty.rs

1use anyhow::{Context as _, Result};
2use base64::{Engine, engine};
3use image::{DynamicImage, imageops::FilterType};
4
5use rustix::event::{PollFd, PollFlags, Timespec, poll};
6use rustix::io::read;
7use rustix::termios::{LocalModes, OptionalActions, tcgetattr, tcgetwinsize, tcsetattr};
8
9use std::io::{Write, stdout};
10use std::os::fd::AsFd as _;
11use std::time::Instant;
12
13pub struct KittyBackend;
14
15impl KittyBackend {
16    pub fn supported() -> Result<bool> {
17        let stdin = std::io::stdin();
18        // save terminal attributes and disable canonical input processing mode
19        let old_attributes = {
20            let old = tcgetattr(&stdin).context("Failed to recieve terminal attibutes")?;
21
22            let mut new = old.clone();
23            new.local_modes &= !LocalModes::ICANON;
24            new.local_modes &= !LocalModes::ECHO;
25            tcsetattr(&stdin, OptionalActions::Now, &new)
26                .context("Failed to update terminal attributes")?;
27            old
28        };
29
30        // generate red rgba test image
31        let mut test_image = Vec::<u8>::with_capacity(32 * 32 * 4);
32        test_image.extend(std::iter::repeat_n([255, 0, 0, 255].iter(), 32 * 32).flatten());
33
34        // print the test image with the action set to query
35        print!(
36            "\x1B_Gi=1,f=32,s=32,v=32,a=q;{}\x1B\\",
37            engine::general_purpose::STANDARD.encode(&test_image)
38        );
39        stdout().flush()?;
40
41        let start_time = Instant::now();
42        let stdin_fd = stdin.as_fd();
43        let mut stdin_pollfd = [PollFd::new(&stdin_fd, PollFlags::IN)];
44        let allowed_bytes = [0x1B, b'_', b'G', b'\\'];
45        let mut buf = Vec::<u8>::new();
46        loop {
47            // check for timeout while polling to avoid blocking the main thread
48            while poll(&mut stdin_pollfd, Some(&Timespec::default()))? < 1 {
49                if start_time.elapsed().as_millis() > 50 {
50                    tcsetattr(&stdin, OptionalActions::Now, &old_attributes)
51                        .context("Failed to update terminal attributes")?;
52                    return Ok(false);
53                }
54            }
55            let mut byte = [0];
56            read(&stdin, &mut byte)?;
57            if allowed_bytes.contains(&byte[0]) {
58                buf.push(byte[0]);
59            }
60            if buf.starts_with(&[0x1B, b'_', b'G']) && buf.ends_with(&[0x1B, b'\\']) {
61                tcsetattr(&stdin, OptionalActions::Now, &old_attributes)
62                    .context("Failed to update terminal attributes")?;
63                return Ok(true);
64            }
65        }
66    }
67}
68
69impl super::ImageBackend for KittyBackend {
70    fn add_image(
71        &self,
72        lines: Vec<String>,
73        image: &DynamicImage,
74        _colors: usize,
75    ) -> Result<String> {
76        let tty_size = tcgetwinsize(std::io::stdin())?;
77        let width_ratio = f64::from(tty_size.ws_col) / f64::from(tty_size.ws_xpixel);
78        let height_ratio = f64::from(tty_size.ws_row) / f64::from(tty_size.ws_ypixel);
79
80        // resize image to fit the text height with the Lanczos3 algorithm
81        let image = image.resize(
82            u32::MAX,
83            (lines.len() as f64 / height_ratio) as u32,
84            FilterType::Lanczos3,
85        );
86        let _image_columns = width_ratio * f64::from(image.width());
87        let image_rows = height_ratio * f64::from(image.height());
88
89        // convert the image to rgba samples
90        let rgba_image = image.to_rgba8();
91        let flat_samples = rgba_image.as_flat_samples();
92        let raw_image = flat_samples
93            .image_slice()
94            .expect("Conversion from image to rgba samples failed");
95        assert_eq!(
96            image.width() as usize * image.height() as usize * 4,
97            raw_image.len()
98        );
99
100        let encoded_image = engine::general_purpose::STANDARD.encode(raw_image); // image data is base64 encoded
101        let mut image_data = Vec::<u8>::new();
102        for chunk in encoded_image.as_bytes().chunks(4096) {
103            // send a 4096 byte chunk of base64 encoded rgba image data
104            image_data.extend(
105                format!(
106                    "\x1B_Gf=32,s={},v={},m=1,a=T;",
107                    image.width(),
108                    image.height()
109                )
110                .as_bytes(),
111            );
112            image_data.extend(chunk);
113            image_data.extend(b"\x1B\\");
114        }
115        image_data.extend(b"\x1B_Gm=0;\x1B\\"); // write empty last chunk
116        image_data.extend(format!("\x1B[{}A", image_rows as u32 - 1).as_bytes()); // move cursor to start of image
117        let mut i = 0;
118        for line in &lines {
119            image_data.extend(format!("\x1B[s{line}\x1B[u\x1B[1B").as_bytes());
120            i += 1;
121        }
122        image_data
123            .extend(format!("\n\x1B[{}B", lines.len().max(image_rows as usize) - i).as_bytes()); // move cursor to end of image
124
125        Ok(String::from_utf8(image_data)?)
126    }
127}