margin.go

  1package ansi
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"io"
  7	"strings"
  8
  9	"github.com/charmbracelet/x/ansi"
 10	"github.com/charmbracelet/x/cellbuf"
 11)
 12
 13// MarginWriter is a Writer that applies indentation and padding around
 14// whatever you write to it.
 15type MarginWriter struct {
 16	w  io.Writer
 17	iw *IndentWriter
 18}
 19
 20// NewMarginWriter returns a new MarginWriter.
 21func NewMarginWriter(ctx RenderContext, w io.Writer, rules StyleBlock) *MarginWriter {
 22	bs := ctx.blockStack
 23
 24	var indentation uint
 25	var margin uint
 26	if rules.Indent != nil {
 27		indentation = *rules.Indent
 28	}
 29	if rules.Margin != nil {
 30		margin = *rules.Margin
 31	}
 32
 33	pw := NewPaddingWriter(w, int(bs.Width(ctx)), func(_ io.Writer) { //nolint:gosec
 34		_, _ = renderText(w, rules.StylePrimitive, " ")
 35	})
 36
 37	ic := " "
 38	if rules.IndentToken != nil {
 39		ic = *rules.IndentToken
 40	}
 41	iw := NewIndentWriter(pw, int(indentation+margin), func(_ io.Writer) { //nolint:gosec
 42		_, _ = renderText(w, bs.Parent().Style.StylePrimitive, ic)
 43	})
 44
 45	return &MarginWriter{
 46		w:  cellbuf.NewPenWriter(w),
 47		iw: iw,
 48	}
 49}
 50
 51func (w *MarginWriter) Write(b []byte) (int, error) {
 52	n, err := w.iw.Write(b)
 53	if err != nil {
 54		return 0, fmt.Errorf("glamour: error writing bytes: %w", err)
 55	}
 56	return n, nil
 57}
 58
 59// PaddingFunc is a function that applies padding around whatever you write to it.
 60type PaddingFunc = func(w io.Writer)
 61
 62// PaddingWriter is a writer that applies padding around whatever you write to
 63// it.
 64type PaddingWriter struct {
 65	Padding int
 66	PadFunc PaddingFunc
 67	w       *cellbuf.PenWriter
 68	cache   bytes.Buffer
 69}
 70
 71// NewPaddingWriter returns a new PaddingWriter.
 72func NewPaddingWriter(w io.Writer, padding int, padFunc PaddingFunc) *PaddingWriter {
 73	return &PaddingWriter{
 74		Padding: padding,
 75		PadFunc: padFunc,
 76		w:       cellbuf.NewPenWriter(w),
 77	}
 78}
 79
 80// Write writes to the padding writer.
 81func (w *PaddingWriter) Write(p []byte) (int, error) {
 82	for i := 0; i < len(p); i++ {
 83		if p[i] == '\n' { //nolint:nestif
 84			line := w.cache.String()
 85			linew := ansi.StringWidth(line)
 86			if w.Padding > 0 && linew < w.Padding {
 87				if w.PadFunc != nil {
 88					for n := 0; n < w.Padding-linew; n++ {
 89						w.PadFunc(w.w)
 90					}
 91				} else {
 92					_, err := io.WriteString(w.w, strings.Repeat(" ", w.Padding-linew))
 93					if err != nil {
 94						return 0, fmt.Errorf("glamour: error writing padding: %w", err)
 95					}
 96				}
 97			}
 98			w.cache.Reset()
 99		} else {
100			w.cache.WriteByte(p[i])
101		}
102
103		_, err := w.w.Write(p[i : i+1])
104		if err != nil {
105			return 0, fmt.Errorf("glamour: error writing bytes: %w", err)
106		}
107	}
108
109	return len(p), nil
110}
111
112// IndentFunc is a function that applies indentation around whatever you write to
113// it.
114type IndentFunc = func(w io.Writer)
115
116// IndentWriter is a writer that applies indentation around whatever you write to
117// it.
118type IndentWriter struct {
119	Indent     int
120	IndentFunc PaddingFunc
121	w          io.Writer
122	pw         *cellbuf.PenWriter
123	skipIndent bool
124}
125
126// NewIndentWriter returns a new IndentWriter.
127func NewIndentWriter(w io.Writer, indent int, indentFunc IndentFunc) *IndentWriter {
128	return &IndentWriter{
129		Indent:     indent,
130		IndentFunc: indentFunc,
131		pw:         cellbuf.NewPenWriter(w),
132		w:          w,
133	}
134}
135
136func (w *IndentWriter) resetPen() {
137	style := w.pw.Style()
138	link := w.pw.Link()
139	if !style.Empty() {
140		_, _ = io.WriteString(w.w, ansi.ResetStyle)
141	}
142	if !link.Empty() {
143		_, _ = io.WriteString(w.w, ansi.ResetHyperlink())
144	}
145}
146
147func (w *IndentWriter) restorePen() {
148	style := w.pw.Style()
149	link := w.pw.Link()
150	if !style.Empty() {
151		_, _ = io.WriteString(w.w, style.Sequence())
152	}
153	if !link.Empty() {
154		_, _ = io.WriteString(w.w, ansi.SetHyperlink(link.URL, link.Params))
155	}
156}
157
158// Write writes to the indentation writer.
159func (w *IndentWriter) Write(p []byte) (int, error) {
160	for i := 0; i < len(p); i++ {
161		if !w.skipIndent {
162			w.resetPen()
163			if w.IndentFunc != nil {
164				for i := 0; i < w.Indent; i++ {
165					w.IndentFunc(w.pw)
166				}
167			} else {
168				_, err := io.WriteString(w.pw, strings.Repeat(" ", w.Indent))
169				if err != nil {
170					return 0, fmt.Errorf("glamour: error writing indentation: %w", err)
171				}
172			}
173
174			w.skipIndent = true
175			w.restorePen()
176		}
177
178		if p[i] == '\n' {
179			w.skipIndent = false
180		}
181
182		_, err := w.pw.Write(p[i : i+1])
183		if err != nil {
184			return 0, fmt.Errorf("glamour: error writing bytes: %w", err)
185		}
186	}
187
188	return len(p), nil
189}