1package hclog
2
3import (
4	"bytes"
5	"io"
6)
7
8type writer struct {
9	b     bytes.Buffer
10	w     io.Writer
11	color ColorOption
12}
13
14func newWriter(w io.Writer, color ColorOption) *writer {
15	return &writer{w: w, color: color}
16}
17
18func (w *writer) Flush(level Level) (err error) {
19	var unwritten = w.b.Bytes()
20
21	if w.color != ColorOff {
22		color := _levelToColor[level]
23		unwritten = []byte(color.Sprintf("%s", unwritten))
24	}
25
26	if lw, ok := w.w.(LevelWriter); ok {
27		_, err = lw.LevelWrite(level, unwritten)
28	} else {
29		_, err = w.w.Write(unwritten)
30	}
31	w.b.Reset()
32	return err
33}
34
35func (w *writer) Write(p []byte) (int, error) {
36	return w.b.Write(p)
37}
38
39func (w *writer) WriteByte(c byte) error {
40	return w.b.WriteByte(c)
41}
42
43func (w *writer) WriteString(s string) (int, error) {
44	return w.b.WriteString(s)
45}
46
47// LevelWriter is the interface that wraps the LevelWrite method.
48type LevelWriter interface {
49	LevelWrite(level Level, p []byte) (n int, err error)
50}
51
52// LeveledWriter writes all log messages to the standard writer,
53// except for log levels that are defined in the overrides map.
54type LeveledWriter struct {
55	standard  io.Writer
56	overrides map[Level]io.Writer
57}
58
59// NewLeveledWriter returns an initialized LeveledWriter.
60//
61// standard will be used as the default writer for all log levels,
62// except for log levels that are defined in the overrides map.
63func NewLeveledWriter(standard io.Writer, overrides map[Level]io.Writer) *LeveledWriter {
64	return &LeveledWriter{
65		standard:  standard,
66		overrides: overrides,
67	}
68}
69
70// Write implements io.Writer.
71func (lw *LeveledWriter) Write(p []byte) (int, error) {
72	return lw.standard.Write(p)
73}
74
75// LevelWrite implements LevelWriter.
76func (lw *LeveledWriter) LevelWrite(level Level, p []byte) (int, error) {
77	w, ok := lw.overrides[level]
78	if !ok {
79		w = lw.standard
80	}
81	return w.Write(p)
82}
83