ansi: kitty graphics: implement basic protocol features

ansi/graphics.go

@@ -0,0 +1,199 @@
+package ansi
+
+import (
+	"bytes"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"image"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/charmbracelet/x/ansi/kitty"
+)
+
+// KittyGraphics returns a sequence that encodes the given image in the Kitty
+// graphics protocol.
+//
+//	APC G [comma separated options] ; [base64 encoded payload] ST
+//
+// See https://sw.kovidgoyal.net/kitty/graphics-protocol/
+func KittyGraphics(payload []byte, opts ...string) string {
+	var buf bytes.Buffer
+	buf.WriteString("\x1b_G")
+	buf.WriteString(strings.Join(opts, ","))
+	if len(payload) > 0 {
+		buf.WriteString(";")
+		buf.Write(payload)
+	}
+	buf.WriteString("\x1b\\")
+	return buf.String()
+}
+
+var (
+	// KittyGraphicsTempDir is the directory where temporary files are stored.
+	// This is used in [WriteKittyGraphics] along with [os.CreateTemp].
+	KittyGraphicsTempDir = ""
+
+	// KittyGraphicsTempPattern is the pattern used to create temporary files.
+	// This is used in [WriteKittyGraphics] along with [os.CreateTemp].
+	// The Kitty Graphics protocol requires the file path to contain the
+	// substring "tty-graphics-protocol".
+	KittyGraphicsTempPattern = "tty-graphics-protocol-*"
+)
+
+// WriteKittyGraphics writes an image using the Kitty Graphics protocol with
+// the given options to w. It chunks the written data if o.Chunk is true.
+//
+// You can omit m and use nil when rendering an image from a file. In this
+// case, you must provide a file path in o.File and use o.Transmission =
+// [kitty.File]. You can also use o.Transmission = [kitty.TempFile] to write
+// the image to a temporary file. In that case, the file path is ignored, and
+// the image is written to a temporary file that is automatically deleted by
+// the terminal.
+//
+// See https://sw.kovidgoyal.net/kitty/graphics-protocol/
+func WriteKittyGraphics(w io.Writer, m image.Image, o *kitty.Options) error {
+	if o == nil {
+		o = &kitty.Options{}
+	}
+
+	if o.Transmission == 0 && len(o.File) != 0 {
+		o.Transmission = kitty.File
+	}
+
+	var data bytes.Buffer // the data to be encoded into base64
+	e := &kitty.Encoder{
+		Compress: o.Compression == kitty.Zlib,
+		Format:   o.Format,
+	}
+
+	switch o.Transmission {
+	case kitty.Direct:
+		if err := e.Encode(&data, m); err != nil {
+			return fmt.Errorf("failed to encode direct image: %w", err)
+		}
+
+	case kitty.SharedMemory:
+		// TODO: Implement shared memory
+		return fmt.Errorf("shared memory transmission is not yet implemented")
+
+	case kitty.File:
+		if len(o.File) == 0 {
+			return kitty.ErrMissingFile
+		}
+
+		f, err := os.Open(o.File)
+		if err != nil {
+			return fmt.Errorf("failed to open file: %w", err)
+		}
+
+		defer f.Close() //nolint:errcheck
+
+		stat, err := f.Stat()
+		if err != nil {
+			return fmt.Errorf("failed to get file info: %w", err)
+		}
+
+		mode := stat.Mode()
+		if !mode.IsRegular() {
+			return fmt.Errorf("file is not a regular file")
+		}
+
+		// Write the file path to the buffer
+		if _, err := data.WriteString(f.Name()); err != nil {
+			return fmt.Errorf("failed to write file path to buffer: %w", err)
+		}
+
+	case kitty.TempFile:
+		f, err := os.CreateTemp(KittyGraphicsTempDir, KittyGraphicsTempPattern)
+		if err != nil {
+			return fmt.Errorf("failed to create file: %w", err)
+		}
+
+		defer f.Close() //nolint:errcheck
+
+		if err := e.Encode(f, m); err != nil {
+			return fmt.Errorf("failed to encode image to file: %w", err)
+		}
+
+		// Write the file path to the buffer
+		if _, err := data.WriteString(f.Name()); err != nil {
+			return fmt.Errorf("failed to write file path to buffer: %w", err)
+		}
+	}
+
+	// Encode image to base64
+	var payload bytes.Buffer // the base64 encoded image to be written to w
+	b64 := base64.NewEncoder(base64.StdEncoding, &payload)
+	if _, err := data.WriteTo(b64); err != nil {
+		return fmt.Errorf("failed to write base64 encoded image to payload: %w", err)
+	}
+	if err := b64.Close(); err != nil {
+		return err
+	}
+
+	// If not chunking, write all at once
+	if !o.Chunk {
+		_, err := io.WriteString(w, KittyGraphics(payload.Bytes(), o.Options()...))
+		return err
+	}
+
+	// Write in chunks
+	var (
+		err error
+		n   int
+	)
+	chunk := make([]byte, kitty.MaxChunkSize)
+	isFirstChunk := true
+
+	for {
+		// Stop if we read less than the chunk size [kitty.MaxChunkSize].
+		n, err = io.ReadFull(&payload, chunk)
+		if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return fmt.Errorf("failed to read chunk: %w", err)
+		}
+
+		opts := buildChunkOptions(o, isFirstChunk, false)
+		if _, err := io.WriteString(w, KittyGraphics(chunk[:n], opts...)); err != nil {
+			return err
+		}
+
+		isFirstChunk = false
+	}
+
+	// Write the last chunk
+	opts := buildChunkOptions(o, isFirstChunk, true)
+	_, err = io.WriteString(w, KittyGraphics(chunk[:n], opts...))
+	return err
+}
+
+// buildChunkOptions creates the options slice for a chunk
+func buildChunkOptions(o *kitty.Options, isFirstChunk, isLastChunk bool) []string {
+	var opts []string
+	if isFirstChunk {
+		opts = o.Options()
+	} else {
+		// These options are allowed in subsequent chunks
+		if o.Quite > 0 {
+			opts = append(opts, fmt.Sprintf("q=%d", o.Quite))
+		}
+		if o.Action == kitty.Frame {
+			opts = append(opts, "a=f")
+		}
+	}
+
+	if !isFirstChunk || !isLastChunk {
+		// We don't need to encode the (m=) option when we only have one chunk.
+		if isLastChunk {
+			opts = append(opts, "m=0")
+		} else {
+			opts = append(opts, "m=1")
+		}
+	}
+	return opts
+}

ansi/kitty/decoder.go

@@ -0,0 +1,85 @@
+package kitty
+
+import (
+	"compress/zlib"
+	"fmt"
+	"image"
+	"image/color"
+	"image/png"
+	"io"
+)
+
+// Decoder is a decoder for the Kitty graphics protocol. It supports decoding
+// images in the 24-bit [RGB], 32-bit [RGBA], and [PNG] formats. It can also
+// decompress data using zlib.
+// The default format is 32-bit [RGBA].
+type Decoder struct {
+	// Uses zlib decompression.
+	Decompress bool
+
+	// Can be one of [RGB], [RGBA], or [PNG].
+	Format int
+
+	// Width of the image in pixels. This can be omitted if the image is [PNG]
+	// formatted.
+	Width int
+
+	// Height of the image in pixels. This can be omitted if the image is [PNG]
+	// formatted.
+	Height int
+}
+
+// Decode decodes the image data from r in the specified format.
+func (d *Decoder) Decode(r io.Reader) (image.Image, error) {
+	if d.Decompress {
+		zr, err := zlib.NewReader(r)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create zlib reader: %w", err)
+		}
+
+		defer zr.Close() //nolint:errcheck
+		r = zr
+	}
+
+	if d.Format == 0 {
+		d.Format = RGBA
+	}
+
+	switch d.Format {
+	case RGBA, RGB:
+		return d.decodeRGBA(r, d.Format == RGBA)
+
+	case PNG:
+		return png.Decode(r)
+
+	default:
+		return nil, fmt.Errorf("unsupported format: %d", d.Format)
+	}
+}
+
+// decodeRGBA decodes the image data in 32-bit RGBA or 24-bit RGB formats.
+func (d *Decoder) decodeRGBA(r io.Reader, alpha bool) (image.Image, error) {
+	m := image.NewRGBA(image.Rect(0, 0, d.Width, d.Height))
+
+	var buf []byte
+	if alpha {
+		buf = make([]byte, 4)
+	} else {
+		buf = make([]byte, 3)
+	}
+
+	for y := 0; y < d.Height; y++ {
+		for x := 0; x < d.Width; x++ {
+			if _, err := io.ReadFull(r, buf[:]); err != nil {
+				return nil, fmt.Errorf("failed to read pixel data: %w", err)
+			}
+			if alpha {
+				m.SetRGBA(x, y, color.RGBA{buf[0], buf[1], buf[2], buf[3]})
+			} else {
+				m.SetRGBA(x, y, color.RGBA{buf[0], buf[1], buf[2], 0xff})
+			}
+		}
+	}
+
+	return m, nil
+}

ansi/kitty/encoder.go

@@ -0,0 +1,64 @@
+package kitty
+
+import (
+	"compress/zlib"
+	"fmt"
+	"image"
+	"image/png"
+	"io"
+)
+
+// Encoder is an encoder for the Kitty graphics protocol. It supports encoding
+// images in the 24-bit [RGB], 32-bit [RGBA], and [PNG] formats, and
+// compressing the data using zlib.
+// The default format is 32-bit [RGBA].
+type Encoder struct {
+	// Uses zlib compression.
+	Compress bool
+
+	// Can be one of [RGBA], [RGB], or [PNG].
+	Format int
+}
+
+// Encode encodes the image data in the specified format and writes it to w.
+func (e *Encoder) Encode(w io.Writer, m image.Image) error {
+	if m == nil {
+		return nil
+	}
+
+	if e.Compress {
+		zw := zlib.NewWriter(w)
+		defer zw.Close() //nolint:errcheck
+		w = zw
+	}
+
+	if e.Format == 0 {
+		e.Format = RGBA
+	}
+
+	switch e.Format {
+	case RGBA, RGB:
+		bounds := m.Bounds()
+		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
+			for x := bounds.Min.X; x < bounds.Max.X; x++ {
+				r, g, b, a := m.At(x, y).RGBA()
+				switch e.Format {
+				case RGBA:
+					w.Write([]byte{byte(r >> 8), byte(g >> 8), byte(b >> 8), byte(a >> 8)}) //nolint:errcheck
+				case RGB:
+					w.Write([]byte{byte(r >> 8), byte(g >> 8), byte(b >> 8)}) //nolint:errcheck
+				}
+			}
+		}
+
+	case PNG:
+		if err := png.Encode(w, m); err != nil {
+			return fmt.Errorf("failed to encode PNG: %w", err)
+		}
+
+	default:
+		return fmt.Errorf("unsupported format: %d", e.Format)
+	}
+
+	return nil
+}