main
Raw Download raw file
  1package ansi
  2
  3import (
  4	"bytes"
  5	"encoding/base64"
  6	"errors"
  7	"fmt"
  8	"image"
  9	"io"
 10	"os"
 11	"strings"
 12
 13	"github.com/charmbracelet/x/ansi/kitty"
 14)
 15
 16// KittyGraphics returns a sequence that encodes the given image in the Kitty
 17// graphics protocol.
 18//
 19//	APC G [comma separated options] ; [base64 encoded payload] ST
 20//
 21// See https://sw.kovidgoyal.net/kitty/graphics-protocol/
 22func KittyGraphics(payload []byte, opts ...string) string {
 23	var buf bytes.Buffer
 24	buf.WriteString("\x1b_G")
 25	buf.WriteString(strings.Join(opts, ","))
 26	if len(payload) > 0 {
 27		buf.WriteString(";")
 28		buf.Write(payload)
 29	}
 30	buf.WriteString("\x1b\\")
 31	return buf.String()
 32}
 33
 34var (
 35	// KittyGraphicsTempDir is the directory where temporary files are stored.
 36	// This is used in [WriteKittyGraphics] along with [os.CreateTemp].
 37	KittyGraphicsTempDir = ""
 38
 39	// KittyGraphicsTempPattern is the pattern used to create temporary files.
 40	// This is used in [WriteKittyGraphics] along with [os.CreateTemp].
 41	// The Kitty Graphics protocol requires the file path to contain the
 42	// substring "tty-graphics-protocol".
 43	KittyGraphicsTempPattern = "tty-graphics-protocol-*"
 44)
 45
 46// WriteKittyGraphics writes an image using the Kitty Graphics protocol with
 47// the given options to w. It chunks the written data if o.Chunk is true.
 48//
 49// You can omit m and use nil when rendering an image from a file. In this
 50// case, you must provide a file path in o.File and use o.Transmission =
 51// [kitty.File]. You can also use o.Transmission = [kitty.TempFile] to write
 52// the image to a temporary file. In that case, the file path is ignored, and
 53// the image is written to a temporary file that is automatically deleted by
 54// the terminal.
 55//
 56// See https://sw.kovidgoyal.net/kitty/graphics-protocol/
 57func WriteKittyGraphics(w io.Writer, m image.Image, o *kitty.Options) error {
 58	if o == nil {
 59		o = &kitty.Options{}
 60	}
 61
 62	if o.Transmission == 0 && len(o.File) != 0 {
 63		o.Transmission = kitty.File
 64	}
 65
 66	var data bytes.Buffer // the data to be encoded into base64
 67	e := &kitty.Encoder{
 68		Compress: o.Compression == kitty.Zlib,
 69		Format:   o.Format,
 70	}
 71
 72	switch o.Transmission {
 73	case kitty.Direct:
 74		if err := e.Encode(&data, m); err != nil {
 75			return fmt.Errorf("failed to encode direct image: %w", err)
 76		}
 77
 78	case kitty.SharedMemory:
 79		// TODO: Implement shared memory
 80		return fmt.Errorf("shared memory transmission is not yet implemented")
 81
 82	case kitty.File:
 83		if len(o.File) == 0 {
 84			return kitty.ErrMissingFile
 85		}
 86
 87		f, err := os.Open(o.File)
 88		if err != nil {
 89			return fmt.Errorf("failed to open file: %w", err)
 90		}
 91
 92		defer f.Close() //nolint:errcheck
 93
 94		stat, err := f.Stat()
 95		if err != nil {
 96			return fmt.Errorf("failed to get file info: %w", err)
 97		}
 98
 99		mode := stat.Mode()
100		if !mode.IsRegular() {
101			return fmt.Errorf("file is not a regular file")
102		}
103
104		// Write the file path to the buffer
105		if _, err := data.WriteString(f.Name()); err != nil {
106			return fmt.Errorf("failed to write file path to buffer: %w", err)
107		}
108
109	case kitty.TempFile:
110		f, err := os.CreateTemp(KittyGraphicsTempDir, KittyGraphicsTempPattern)
111		if err != nil {
112			return fmt.Errorf("failed to create file: %w", err)
113		}
114
115		defer f.Close() //nolint:errcheck
116
117		if err := e.Encode(f, m); err != nil {
118			return fmt.Errorf("failed to encode image to file: %w", err)
119		}
120
121		// Write the file path to the buffer
122		if _, err := data.WriteString(f.Name()); err != nil {
123			return fmt.Errorf("failed to write file path to buffer: %w", err)
124		}
125	}
126
127	// Encode image to base64
128	var payload bytes.Buffer // the base64 encoded image to be written to w
129	b64 := base64.NewEncoder(base64.StdEncoding, &payload)
130	if _, err := data.WriteTo(b64); err != nil {
131		return fmt.Errorf("failed to write base64 encoded image to payload: %w", err)
132	}
133	if err := b64.Close(); err != nil {
134		return err
135	}
136
137	// If not chunking, write all at once
138	if !o.Chunk {
139		_, err := io.WriteString(w, KittyGraphics(payload.Bytes(), o.Options()...))
140		return err
141	}
142
143	// Write in chunks
144	var (
145		err error
146		n   int
147	)
148	chunk := make([]byte, kitty.MaxChunkSize)
149	isFirstChunk := true
150
151	for {
152		// Stop if we read less than the chunk size [kitty.MaxChunkSize].
153		n, err = io.ReadFull(&payload, chunk)
154		if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) {
155			break
156		}
157		if err != nil {
158			return fmt.Errorf("failed to read chunk: %w", err)
159		}
160
161		opts := buildChunkOptions(o, isFirstChunk, false)
162		if _, err := io.WriteString(w, KittyGraphics(chunk[:n], opts...)); err != nil {
163			return err
164		}
165
166		isFirstChunk = false
167	}
168
169	// Write the last chunk
170	opts := buildChunkOptions(o, isFirstChunk, true)
171	_, err = io.WriteString(w, KittyGraphics(chunk[:n], opts...))
172	return err
173}
174
175// buildChunkOptions creates the options slice for a chunk
176func buildChunkOptions(o *kitty.Options, isFirstChunk, isLastChunk bool) []string {
177	var opts []string
178	if isFirstChunk {
179		opts = o.Options()
180	} else {
181		// These options are allowed in subsequent chunks
182		if o.Quite > 0 {
183			opts = append(opts, fmt.Sprintf("q=%d", o.Quite))
184		}
185		if o.Action == kitty.Frame {
186			opts = append(opts, "a=f")
187		}
188	}
189
190	if !isFirstChunk || !isLastChunk {
191		// We don't need to encode the (m=) option when we only have one chunk.
192		if isLastChunk {
193			opts = append(opts, "m=0")
194		} else {
195			opts = append(opts, "m=1")
196		}
197	}
198	return opts
199}