1package text
2
3import (
4	"io"
5)
6
7// Indent inserts prefix at the beginning of each non-empty line of s. The
8// end-of-line marker is NL.
9func Indent(s, prefix string) string {
10	return string(IndentBytes([]byte(s), []byte(prefix)))
11}
12
13// IndentBytes inserts prefix at the beginning of each non-empty line of b.
14// The end-of-line marker is NL.
15func IndentBytes(b, prefix []byte) []byte {
16	var res []byte
17	bol := true
18	for _, c := range b {
19		if bol && c != '\n' {
20			res = append(res, prefix...)
21		}
22		res = append(res, c)
23		bol = c == '\n'
24	}
25	return res
26}
27
28// Writer indents each line of its input.
29type indentWriter struct {
30	w   io.Writer
31	bol bool
32	pre [][]byte
33	sel int
34	off int
35}
36
37// NewIndentWriter makes a new write filter that indents the input
38// lines. Each line is prefixed in order with the corresponding
39// element of pre. If there are more lines than elements, the last
40// element of pre is repeated for each subsequent line.
41func NewIndentWriter(w io.Writer, pre ...[]byte) io.Writer {
42	return &indentWriter{
43		w:   w,
44		pre: pre,
45		bol: true,
46	}
47}
48
49// The only errors returned are from the underlying indentWriter.
50func (w *indentWriter) Write(p []byte) (n int, err error) {
51	for _, c := range p {
52		if w.bol {
53			var i int
54			i, err = w.w.Write(w.pre[w.sel][w.off:])
55			w.off += i
56			if err != nil {
57				return n, err
58			}
59		}
60		_, err = w.w.Write([]byte{c})
61		if err != nil {
62			return n, err
63		}
64		n++
65		w.bol = c == '\n'
66		if w.bol {
67			w.off = 0
68			if w.sel < len(w.pre)-1 {
69				w.sel++
70			}
71		}
72	}
73	return n, nil
74}
75