writer.go

  1package cellbuf
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"strings"
  7
  8	"github.com/charmbracelet/x/ansi"
  9)
 10
 11// CellBuffer is a cell buffer that represents a set of cells in a screen or a
 12// grid.
 13type CellBuffer interface {
 14	// Cell returns the cell at the given position.
 15	Cell(x, y int) *Cell
 16	// SetCell sets the cell at the given position to the given cell. It
 17	// returns whether the cell was set successfully.
 18	SetCell(x, y int, c *Cell) bool
 19	// Bounds returns the bounds of the cell buffer.
 20	Bounds() Rectangle
 21}
 22
 23// FillRect fills the rectangle within the cell buffer with the given cell.
 24// This will not fill cells outside the bounds of the cell buffer.
 25func FillRect(s CellBuffer, c *Cell, rect Rectangle) {
 26	for y := rect.Min.Y; y < rect.Max.Y; y++ {
 27		for x := rect.Min.X; x < rect.Max.X; x++ {
 28			s.SetCell(x, y, c) //nolint:errcheck
 29		}
 30	}
 31}
 32
 33// Fill fills the cell buffer with the given cell.
 34func Fill(s CellBuffer, c *Cell) {
 35	FillRect(s, c, s.Bounds())
 36}
 37
 38// ClearRect clears the rectangle within the cell buffer with blank cells.
 39func ClearRect(s CellBuffer, rect Rectangle) {
 40	FillRect(s, nil, rect)
 41}
 42
 43// Clear clears the cell buffer with blank cells.
 44func Clear(s CellBuffer) {
 45	Fill(s, nil)
 46}
 47
 48// SetContentRect clears the rectangle within the cell buffer with blank cells,
 49// and sets the given string as its content. If the height or width of the
 50// string exceeds the height or width of the cell buffer, it will be truncated.
 51func SetContentRect(s CellBuffer, str string, rect Rectangle) {
 52	// Replace all "\n" with "\r\n" to ensure the cursor is reset to the start
 53	// of the line. Make sure we don't replace "\r\n" with "\r\r\n".
 54	str = strings.ReplaceAll(str, "\r\n", "\n")
 55	str = strings.ReplaceAll(str, "\n", "\r\n")
 56	ClearRect(s, rect)
 57	printString(s, ansi.GraphemeWidth, rect.Min.X, rect.Min.Y, rect, str, true, "")
 58}
 59
 60// SetContent clears the cell buffer with blank cells, and sets the given string
 61// as its content. If the height or width of the string exceeds the height or
 62// width of the cell buffer, it will be truncated.
 63func SetContent(s CellBuffer, str string) {
 64	SetContentRect(s, str, s.Bounds())
 65}
 66
 67// Render returns a string representation of the grid with ANSI escape sequences.
 68func Render(d CellBuffer) string {
 69	var buf bytes.Buffer
 70	height := d.Bounds().Dy()
 71	for y := 0; y < height; y++ {
 72		_, line := RenderLine(d, y)
 73		buf.WriteString(line)
 74		if y < height-1 {
 75			buf.WriteString("\r\n")
 76		}
 77	}
 78	return buf.String()
 79}
 80
 81// RenderLine returns a string representation of the yth line of the grid along
 82// with the width of the line.
 83func RenderLine(d CellBuffer, n int) (w int, line string) {
 84	var pen Style
 85	var link Link
 86	var buf bytes.Buffer
 87	var pendingLine string
 88	var pendingWidth int // this ignores space cells until we hit a non-space cell
 89
 90	writePending := func() {
 91		// If there's no pending line, we don't need to do anything.
 92		if len(pendingLine) == 0 {
 93			return
 94		}
 95		buf.WriteString(pendingLine)
 96		w += pendingWidth
 97		pendingWidth = 0
 98		pendingLine = ""
 99	}
100
101	for x := 0; x < d.Bounds().Dx(); x++ {
102		if cell := d.Cell(x, n); cell != nil && cell.Width > 0 {
103			// Convert the cell's style and link to the given color profile.
104			cellStyle := cell.Style
105			cellLink := cell.Link
106			if cellStyle.Empty() && !pen.Empty() {
107				writePending()
108				buf.WriteString(ansi.ResetStyle) //nolint:errcheck
109				pen.Reset()
110			}
111			if !cellStyle.Equal(&pen) {
112				writePending()
113				seq := cellStyle.DiffSequence(pen)
114				buf.WriteString(seq) // nolint:errcheck
115				pen = cellStyle
116			}
117
118			// Write the URL escape sequence
119			if cellLink != link && link.URL != "" {
120				writePending()
121				buf.WriteString(ansi.ResetHyperlink()) //nolint:errcheck
122				link.Reset()
123			}
124			if cellLink != link {
125				writePending()
126				buf.WriteString(ansi.SetHyperlink(cellLink.URL, cellLink.Params)) //nolint:errcheck
127				link = cellLink
128			}
129
130			// We only write the cell content if it's not empty. If it is, we
131			// append it to the pending line and width to be evaluated later.
132			if cell.Equal(&BlankCell) {
133				pendingLine += cell.String()
134				pendingWidth += cell.Width
135			} else {
136				writePending()
137				buf.WriteString(cell.String())
138				w += cell.Width
139			}
140		}
141	}
142	if link.URL != "" {
143		buf.WriteString(ansi.ResetHyperlink()) //nolint:errcheck
144	}
145	if !pen.Empty() {
146		buf.WriteString(ansi.ResetStyle) //nolint:errcheck
147	}
148	return w, strings.TrimRight(buf.String(), " ") // Trim trailing spaces
149}
150
151// ScreenWriter represents a writer that writes to a [Screen] parsing ANSI
152// escape sequences and Unicode characters and converting them into cells that
153// can be written to a cell [Buffer].
154type ScreenWriter struct {
155	*Screen
156}
157
158// NewScreenWriter creates a new ScreenWriter that writes to the given Screen.
159// This is a convenience function for creating a ScreenWriter.
160func NewScreenWriter(s *Screen) *ScreenWriter {
161	return &ScreenWriter{s}
162}
163
164// Write writes the given bytes to the screen.
165// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
166// sequences.
167func (s *ScreenWriter) Write(p []byte) (n int, err error) {
168	printString(s.Screen, s.method,
169		s.cur.X, s.cur.Y, s.Bounds(),
170		p, false, "")
171	return len(p), nil
172}
173
174// SetContent clears the screen with blank cells, and sets the given string as
175// its content. If the height or width of the string exceeds the height or
176// width of the screen, it will be truncated.
177//
178// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape sequences.
179func (s *ScreenWriter) SetContent(str string) {
180	s.SetContentRect(str, s.Bounds())
181}
182
183// SetContentRect clears the rectangle within the screen with blank cells, and
184// sets the given string as its content. If the height or width of the string
185// exceeds the height or width of the screen, it will be truncated.
186//
187// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
188// sequences.
189func (s *ScreenWriter) SetContentRect(str string, rect Rectangle) {
190	// Replace all "\n" with "\r\n" to ensure the cursor is reset to the start
191	// of the line. Make sure we don't replace "\r\n" with "\r\r\n".
192	str = strings.ReplaceAll(str, "\r\n", "\n")
193	str = strings.ReplaceAll(str, "\n", "\r\n")
194	s.ClearRect(rect)
195	printString(s.Screen, s.method,
196		rect.Min.X, rect.Min.Y, rect,
197		str, true, "")
198}
199
200// Print prints the string at the current cursor position. It will wrap the
201// string to the width of the screen if it exceeds the width of the screen.
202// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
203// sequences.
204func (s *ScreenWriter) Print(str string, v ...interface{}) {
205	if len(v) > 0 {
206		str = fmt.Sprintf(str, v...)
207	}
208	printString(s.Screen, s.method,
209		s.cur.X, s.cur.Y, s.Bounds(),
210		str, false, "")
211}
212
213// PrintAt prints the string at the given position. It will wrap the string to
214// the width of the screen if it exceeds the width of the screen.
215// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
216// sequences.
217func (s *ScreenWriter) PrintAt(x, y int, str string, v ...interface{}) {
218	if len(v) > 0 {
219		str = fmt.Sprintf(str, v...)
220	}
221	printString(s.Screen, s.method,
222		x, y, s.Bounds(),
223		str, false, "")
224}
225
226// PrintCrop prints the string at the current cursor position and truncates the
227// text if it exceeds the width of the screen. Use tail to specify a string to
228// append if the string is truncated.
229// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
230// sequences.
231func (s *ScreenWriter) PrintCrop(str string, tail string) {
232	printString(s.Screen, s.method,
233		s.cur.X, s.cur.Y, s.Bounds(),
234		str, true, tail)
235}
236
237// PrintCropAt prints the string at the given position and truncates the text
238// if it exceeds the width of the screen. Use tail to specify a string to append
239// if the string is truncated.
240// This will recognize ANSI [ansi.SGR] style and [ansi.SetHyperlink] escape
241// sequences.
242func (s *ScreenWriter) PrintCropAt(x, y int, str string, tail string) {
243	printString(s.Screen, s.method,
244		x, y, s.Bounds(),
245		str, true, tail)
246}
247
248// printString draws a string starting at the given position.
249func printString[T []byte | string](
250	s CellBuffer,
251	m ansi.Method,
252	x, y int,
253	bounds Rectangle, str T,
254	truncate bool, tail string,
255) {
256	p := ansi.GetParser()
257	defer ansi.PutParser(p)
258
259	var tailc Cell
260	if truncate && len(tail) > 0 {
261		if m == ansi.WcWidth {
262			tailc = *NewCellString(tail)
263		} else {
264			tailc = *NewGraphemeCell(tail)
265		}
266	}
267
268	decoder := ansi.DecodeSequenceWc[T]
269	if m == ansi.GraphemeWidth {
270		decoder = ansi.DecodeSequence[T]
271	}
272
273	var cell Cell
274	var style Style
275	var link Link
276	var state byte
277	for len(str) > 0 {
278		seq, width, n, newState := decoder(str, state, p)
279
280		switch width {
281		case 1, 2, 3, 4: // wide cells can go up to 4 cells wide
282			cell.Width += width
283			cell.Append([]rune(string(seq))...)
284
285			if !truncate && x+cell.Width > bounds.Max.X && y+1 < bounds.Max.Y {
286				// Wrap the string to the width of the window
287				x = bounds.Min.X
288				y++
289			}
290			if Pos(x, y).In(bounds) {
291				if truncate && tailc.Width > 0 && x+cell.Width > bounds.Max.X-tailc.Width {
292					// Truncate the string and append the tail if any.
293					cell := tailc
294					cell.Style = style
295					cell.Link = link
296					s.SetCell(x, y, &cell)
297					x += tailc.Width
298				} else {
299					// Print the cell to the screen
300					cell.Style = style
301					cell.Link = link
302					s.SetCell(x, y, &cell) //nolint:errcheck
303					x += width
304				}
305			}
306
307			// String is too long for the line, truncate it.
308			// Make sure we reset the cell for the next iteration.
309			cell.Reset()
310		default:
311			// Valid sequences always have a non-zero Cmd.
312			// TODO: Handle cursor movement and other sequences
313			switch {
314			case ansi.HasCsiPrefix(seq) && p.Command() == 'm':
315				// SGR - Select Graphic Rendition
316				ReadStyle(p.Params(), &style)
317			case ansi.HasOscPrefix(seq) && p.Command() == 8:
318				// Hyperlinks
319				ReadLink(p.Data(), &link)
320			case ansi.Equal(seq, T("\n")):
321				y++
322			case ansi.Equal(seq, T("\r")):
323				x = bounds.Min.X
324			default:
325				cell.Append([]rune(string(seq))...)
326			}
327		}
328
329		// Advance the state and data
330		state = newState
331		str = str[n:]
332	}
333
334	// Make sure to set the last cell if it's not empty.
335	if !cell.Empty() {
336		s.SetCell(x, y, &cell) //nolint:errcheck
337		cell.Reset()
338	}
339}