ansi: kitty graphics: implement basic protocol features (#322)

* feat(ansi): kitty graphics: implement basic protocol features

This adds support for writing images in the Kitty Graphics protocol. The
protocol is used to transmit images over the terminal using escape
sequences. The protocol supports 24-bit RGB, 32-bit RGBA, and PNG
formats. Images can be compressed using zlib and transmitted in chunks.

* feat(ansi): kitty graphics: support unicode placeholders

* fix(ansi): kitty: options validation

* fix(ansi): kitty: closing file

* feat(ansi): kitty: support encoding.TextMarshaler and encoding.TextUnmarshaler

* chore: exclude tests

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
+}