1package term
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"sync"
8
9	"github.com/go-kit/kit/log"
10)
11
12// Color represents an ANSI color. The zero value is Default.
13type Color uint8
14
15// ANSI colors.
16const (
17	Default = Color(iota)
18
19	Black
20	DarkRed
21	DarkGreen
22	Brown
23	DarkBlue
24	DarkMagenta
25	DarkCyan
26	Gray
27
28	DarkGray
29	Red
30	Green
31	Yellow
32	Blue
33	Magenta
34	Cyan
35	White
36
37	numColors
38)
39
40// For more on ANSI escape codes see
41// https://en.wikipedia.org/wiki/ANSI_escape_code. See in particular
42// https://en.wikipedia.org/wiki/ANSI_escape_code#Colors.
43
44var (
45	resetColorBytes = []byte("\x1b[39;49;22m")
46	fgColorBytes    [][]byte
47	bgColorBytes    [][]byte
48)
49
50func init() {
51	// Default
52	fgColorBytes = append(fgColorBytes, []byte("\x1b[39m"))
53	bgColorBytes = append(bgColorBytes, []byte("\x1b[49m"))
54
55	// dark colors
56	for color := Black; color < DarkGray; color++ {
57		fgColorBytes = append(fgColorBytes, []byte(fmt.Sprintf("\x1b[%dm", 30+color-Black)))
58		bgColorBytes = append(bgColorBytes, []byte(fmt.Sprintf("\x1b[%dm", 40+color-Black)))
59	}
60
61	// bright colors
62	for color := DarkGray; color < numColors; color++ {
63		fgColorBytes = append(fgColorBytes, []byte(fmt.Sprintf("\x1b[%d;1m", 30+color-DarkGray)))
64		bgColorBytes = append(bgColorBytes, []byte(fmt.Sprintf("\x1b[%d;1m", 40+color-DarkGray)))
65	}
66}
67
68// FgBgColor represents a foreground and background color.
69type FgBgColor struct {
70	Fg, Bg Color
71}
72
73func (c FgBgColor) isZero() bool {
74	return c.Fg == Default && c.Bg == Default
75}
76
77// NewColorLogger returns a Logger which writes colored logs to w. ANSI color
78// codes for the colors returned by color are added to the formatted output
79// from the Logger returned by newLogger and the combined result written to w.
80func NewColorLogger(w io.Writer, newLogger func(io.Writer) log.Logger, color func(keyvals ...interface{}) FgBgColor) log.Logger {
81	if color == nil {
82		panic("color func nil")
83	}
84	return &colorLogger{
85		w:             w,
86		newLogger:     newLogger,
87		color:         color,
88		bufPool:       sync.Pool{New: func() interface{} { return &loggerBuf{} }},
89		noColorLogger: newLogger(w),
90	}
91}
92
93type colorLogger struct {
94	w             io.Writer
95	newLogger     func(io.Writer) log.Logger
96	color         func(keyvals ...interface{}) FgBgColor
97	bufPool       sync.Pool
98	noColorLogger log.Logger
99}
100
101func (l *colorLogger) Log(keyvals ...interface{}) error {
102	color := l.color(keyvals...)
103	if color.isZero() {
104		return l.noColorLogger.Log(keyvals...)
105	}
106
107	lb := l.getLoggerBuf()
108	defer l.putLoggerBuf(lb)
109	if color.Fg != Default {
110		lb.buf.Write(fgColorBytes[color.Fg])
111	}
112	if color.Bg != Default {
113		lb.buf.Write(bgColorBytes[color.Bg])
114	}
115	err := lb.logger.Log(keyvals...)
116	if err != nil {
117		return err
118	}
119	if color.Fg != Default || color.Bg != Default {
120		lb.buf.Write(resetColorBytes)
121	}
122	_, err = io.Copy(l.w, lb.buf)
123	return err
124}
125
126type loggerBuf struct {
127	buf    *bytes.Buffer
128	logger log.Logger
129}
130
131func (l *colorLogger) getLoggerBuf() *loggerBuf {
132	lb := l.bufPool.Get().(*loggerBuf)
133	if lb.buf == nil {
134		lb.buf = &bytes.Buffer{}
135		lb.logger = l.newLogger(lb.buf)
136	} else {
137		lb.buf.Reset()
138	}
139	return lb
140}
141
142func (l *colorLogger) putLoggerBuf(cb *loggerBuf) {
143	l.bufPool.Put(cb)
144}
145