graphics.go

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