Skip to main content

pulldown_cmark_mdcat/terminal/capabilities/
kitty.rs

1// Copyright 2020 Sebastian Wiesner <sebastian@swsnr.de>
2// Copyright 2019 Fabian Spillner <fabian.spillner@gmail.com>
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7
8//  http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Kitty terminal extensions.
17use std::fmt::Display;
18#[cfg(not(feature = "image-processing"))]
19use std::io::ErrorKind;
20use std::io::{Error, Write};
21use std::str;
22
23use base64::engine::general_purpose::STANDARD;
24use base64::Engine;
25use tracing::{event, instrument, Level};
26
27use crate::resources::image::*;
28use crate::resources::MimeData;
29use crate::terminal::size::{PixelSize, TerminalSize};
30
31/// An error which occurred while rendering or writing an image with the Kitty image protocol.
32#[derive(Debug)]
33pub enum KittyImageError {
34    /// A general IO error.
35    IoError(std::io::Error),
36    /// Processing a pixel image, e.g. for format conversion, failed
37    #[cfg(feature = "image-processing")]
38    ImageError(image::ImageError),
39}
40
41impl Display for KittyImageError {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        match self {
44            KittyImageError::IoError(error) => write!(f, "Failed to render kitty image: {error}"),
45            #[cfg(feature = "image-processing")]
46            KittyImageError::ImageError(image_error) => {
47                write!(f, "Failed to process pixel image: {image_error}")
48            }
49        }
50    }
51}
52
53impl std::error::Error for KittyImageError {
54    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
55        match self {
56            KittyImageError::IoError(error) => Some(error),
57            #[cfg(feature = "image-processing")]
58            KittyImageError::ImageError(image_error) => Some(image_error),
59        }
60    }
61}
62
63impl From<KittyImageError> for std::io::Error {
64    fn from(value: KittyImageError) -> Self {
65        std::io::Error::other(value)
66    }
67}
68
69impl From<std::io::Error> for KittyImageError {
70    fn from(value: std::io::Error) -> Self {
71        Self::IoError(value)
72    }
73}
74
75#[cfg(feature = "image-processing")]
76impl From<image::ImageError> for KittyImageError {
77    fn from(value: image::ImageError) -> Self {
78        Self::ImageError(value)
79    }
80}
81
82/// Image data for the kitty graphics protocol.
83///
84/// See [Terminal graphics protocol][1] for a complete documentation.
85///
86/// [1]: https://sw.kovidgoyal.net/kitty/graphics-protocol/
87enum KittyImageData {
88    Png {
89        data: Vec<u8>,
90        pixel_size: Option<PixelSize>,
91    },
92    #[cfg(feature = "image-processing")]
93    Rgb(PixelSize, Vec<u8>),
94    #[cfg(feature = "image-processing")]
95    Rgba(PixelSize, Vec<u8>),
96}
97
98impl KittyImageData {
99    /// Return the format code for this data for the `f` control data field.
100    ///
101    /// See the [Transferring pixel data][1] for reference.
102    ///
103    /// [1]: https://sw.kovidgoyal.net/kitty/graphics-protocol.html#transferring-pixel-data
104    fn f_format_code(&self) -> &str {
105        match self {
106            KittyImageData::Png { .. } => "100",
107            #[cfg(feature = "image-processing")]
108            KittyImageData::Rgb(_, _) => "24",
109            #[cfg(feature = "image-processing")]
110            KittyImageData::Rgba(_, _) => "32",
111        }
112    }
113
114    /// Get the actual data.
115    fn data(&self) -> &[u8] {
116        match self {
117            KittyImageData::Png { data, .. } => data,
118            #[cfg(feature = "image-processing")]
119            KittyImageData::Rgb(_, ref contents) => contents,
120            #[cfg(feature = "image-processing")]
121            KittyImageData::Rgba(_, ref contents) => contents,
122        }
123    }
124
125    /// Get the size of the image contained in this data.
126    ///
127    /// `Some` if the size is explicitly specified for this data, `None` otherwise, i.e. in PNG
128    /// format).
129    fn size(&self) -> Option<PixelSize> {
130        match self {
131            KittyImageData::Png { pixel_size, .. } => *pixel_size,
132            #[cfg(feature = "image-processing")]
133            KittyImageData::Rgb(size, _) => Some(*size),
134            #[cfg(feature = "image-processing")]
135            KittyImageData::Rgba(size, _) => Some(*size),
136        }
137    }
138}
139
140impl KittyImageData {
141    fn write_to(&self, writer: &mut dyn Write, move_cursor: bool) -> Result<(), Error> {
142        let image_data = STANDARD.encode(self.data());
143        let image_data_chunks = image_data.as_bytes().chunks(4096);
144        let number_of_chunks = image_data_chunks.len();
145
146        for (i, chunk_data) in image_data_chunks.enumerate() {
147            let is_first_chunk = i == 0;
148            // The value for the m field
149            let m = if i < number_of_chunks - 1 { 1 } else { 0 };
150            if is_first_chunk {
151                // For the first chunk we must write the header for the image.
152                //
153                // a=T tells kitty that we transfer image data and want to show the image
154                // immediately.
155                //
156                // t=d tells kitty that we transfer image data inline in the escape code.
157                //
158                // I=1 tells kitty that we want to treat every image as unique and not have kitty
159                // reuse images.  At least wezterm requires this; otherwise past images disappear
160                // because wezterm seems to assume that we're reusing some image ID.
161                //
162                // f tells kitty about the data format.
163                //
164                // m tells kitty whether to expect more chunks or whether this is the last one.
165                //
166                // q=2 tells kitty never to respond to our image sequence; we're not reading these
167                // responses anyway.
168                //
169                let f = self.f_format_code();
170                write!(writer, "\x1b_Ga=T,t=d,I=1,f={f}")?;
171                if let Some(size) = self.size() {
172                    write!(writer, ",s={},v={}", size.x, size.y)?;
173                }
174                if !move_cursor {
175                    write!(writer, ",C=1")?;
176                }
177                write!(writer, ",m={m},q=2;")?;
178            } else {
179                // For follow up chunks we must not repeat the header, but only indicate whether we
180                // expect a response and whether more data is to follow.
181                write!(writer, "\x1b_Gm={m},q=2;")?;
182            }
183            writer.write_all(chunk_data)?;
184            write!(writer, "\x1b\\")?;
185        }
186
187        Ok(())
188    }
189}
190
191/// Provides access to printing images for kitty.
192#[derive(Debug, Copy, Clone)]
193pub struct KittyGraphicsProtocol;
194
195impl KittyGraphicsProtocol {
196    /// Render mime data obtained from `url` and wrap it in a `KittyImage`.
197    ///
198    /// This implementation processes the image to scale it to the given `terminal_size`, and
199    /// supports various pixel image types, as well as SVG.
200    #[cfg(feature = "image-processing")]
201    fn render(
202        self,
203        mime_data: MimeData,
204        terminal_size: TerminalSize,
205    ) -> Result<KittyImageData, KittyImageError> {
206        use image::ImageFormat;
207
208        let image = if let Some("image/svg+xml") = mime_data.mime_type_essence() {
209            event!(Level::DEBUG, "Rendering mime data to SVG");
210            let png_data = crate::resources::svg::render_svg_to_png(&mime_data.data)?;
211            image::load_from_memory_with_format(&png_data, ImageFormat::Png)?
212        } else {
213            let image_format = mime_data
214                .mime_type_essence()
215                .and_then(image::ImageFormat::from_mime_type);
216            match image_format {
217                // If we already have information about the mime type of the resource data let's
218                // use it, and trust whoever provided it to have gotten it right.
219                Some(format) => image::load_from_memory_with_format(&mime_data.data, format)?,
220                // If we don't know the mime type of the original data have image guess the format.
221                None => image::load_from_memory(&mime_data.data)?,
222            }
223        };
224
225        match downsize_to_columns(&image, terminal_size) {
226            Some(downsized_image) => {
227                event!(
228                    Level::DEBUG,
229                    "Image scaled down to column limit, rendering RGB data"
230                );
231                Ok(self.render_as_rgb_or_rgba(downsized_image))
232            }
233            None if mime_data.mime_type_essence() == Some("image/png") => {
234                event!(
235                    Level::DEBUG,
236                    "PNG image of appropriate size, rendering original image data"
237                );
238                Ok(self.render_as_png(mime_data.data))
239            }
240            None => {
241                event!(Level::DEBUG, "Image not in PNG format, rendering RGB data");
242                Ok(self.render_as_rgb_or_rgba(image))
243            }
244        }
245    }
246
247    /// Render mime data obtained from `url` and wrap it in a `KittyImageData`.
248    ///
249    /// This implementation does not support image processing, and only renders PNG images which
250    /// kitty supports directly.
251    #[cfg(not(feature = "image-processing"))]
252    fn render(
253        self,
254        mime_data: MimeData,
255        _terminal_size: TerminalSize,
256    ) -> Result<KittyImageData, KittyImageError> {
257        match mime_data.mime_type_essence() {
258            Some("image/png") => Ok(self.render_as_png(mime_data.data)),
259            _ => {
260                event!(
261                    Level::DEBUG,
262                    "Only PNG images supported without image-processing feature, but got {:?}",
263                    mime_data.mime_type
264                );
265                Err(std::io::Error::new(
266                    ErrorKind::Unsupported,
267                    format!(
268                        "Image data with mime type {:?} not supported",
269                        mime_data.mime_type
270                    ),
271                )
272                .into())
273            }
274        }
275    }
276
277    /// Wrap the image bytes as PNG format in `KittyImage`.
278    fn render_as_png(self, data: Vec<u8>) -> KittyImageData {
279        KittyImageData::Png {
280            data,
281            pixel_size: None,
282        }
283    }
284
285    /// Render the image as RGB/RGBA format and wrap the image bytes in `KittyImage`.
286    ///
287    /// If the image size exceeds `terminal_size` in either dimension scale the
288    /// image down to `terminal_size` (preserving aspect ratio).
289    #[cfg(feature = "image-processing")]
290    fn render_as_rgb_or_rgba(self, image: image::DynamicImage) -> KittyImageData {
291        use image::{ColorType, GenericImageView};
292
293        let size = PixelSize::from_xy(image.dimensions());
294        match image.color() {
295            ColorType::L8 | ColorType::Rgb8 | ColorType::L16 | ColorType::Rgb16 => {
296                KittyImageData::Rgb(size, image.into_rgb8().into_raw())
297            }
298            // Default to RGBA format: We cannot match all colour types because
299            // ColorType is marked non-exhaustive, but RGBA is a safe default
300            // because we can convert any image to RGBA, at worth with additional
301            // runtime costs.
302            _ => KittyImageData::Rgba(size, image.into_rgba8().into_raw()),
303        }
304    }
305}
306
307/// Kitty's inline image protocol.
308///
309/// Kitty's escape sequence is like: Put the command key/value pairs together like "{}={}(,*)"
310/// and write them along with the image bytes in 4096 bytes chunks to the stdout.
311///
312/// Its documentation gives the following python example:
313///
314/// ```python
315/// import sys
316/// from base64 import standard_b64encode
317///
318/// def serialize_gr_command(cmd, payload=None):
319///   cmd = ','.join('{}={}'.format(k, v) for k, v in cmd.items())
320///   ans = []
321///   w = ans.append
322///   w(b'\033_G'), w(cmd.encode('ascii'))
323///   if payload:
324///     w(b';')
325///     w(payload)
326///   w(b'\033\\')
327///   return b''.join(ans)
328///
329/// def write_chunked(cmd, data):
330///   cmd = {'a': 'T', 'f': 100}
331///   data = standard_b64encode(data)
332///   while data:
333///     chunk, data = data[:4096], data[4096:]
334///     m = 1 if data else 0
335///     cmd['m'] = m
336///     sys.stdout.buffer.write(serialize_gr_command(cmd, chunk))
337///     sys.stdout.flush()
338///     cmd.clear()
339/// ```
340///
341/// See <https://sw.kovidgoyal.net/kitty/graphics-protocol.html#control-data-reference>
342/// for reference.
343impl KittyGraphicsProtocol {
344    /// Write raw PNG bytes inline to the terminal.
345    pub(crate) fn write_png_data(
346        &self,
347        writer: &mut dyn Write,
348        png_data: Vec<u8>,
349        move_cursor: bool,
350    ) -> std::io::Result<()> {
351        KittyImageData::Png {
352            data: png_data,
353            pixel_size: None,
354        }
355        .write_to(writer, move_cursor)
356    }
357}
358
359impl InlineImageProtocol for KittyGraphicsProtocol {
360    #[instrument(skip(self, writer, resource_handler, terminal_size))]
361    fn write_inline_image(
362        &self,
363        writer: &mut dyn Write,
364        resource_handler: &dyn crate::ResourceUrlHandler,
365        url: &url::Url,
366        terminal_size: crate::TerminalSize,
367    ) -> std::io::Result<()> {
368        let mime_data = resource_handler.read_resource(url)?;
369        event!(
370            Level::DEBUG,
371            "Received data of mime type {:?}",
372            mime_data.mime_type
373        );
374        let image = self.render(mime_data, terminal_size)?;
375        image.write_to(writer, true)
376    }
377}