github.com/go-kit/log@v0.2.1/term/colorlogger.go (about)

     1  package term
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  
     9  	"github.com/go-kit/log"
    10  )
    11  
    12  // Color represents an ANSI color. The zero value is Default.
    13  type Color uint8
    14  
    15  // ANSI colors.
    16  const (
    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  
    44  var (
    45  	resetColorBytes = []byte("\x1b[39;49;22m")
    46  	fgColorBytes    [][]byte
    47  	bgColorBytes    [][]byte
    48  )
    49  
    50  func 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.
    69  type FgBgColor struct {
    70  	Fg, Bg Color
    71  }
    72  
    73  func (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.
    80  func 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  
    93  type 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  
   101  func (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  
   126  type loggerBuf struct {
   127  	buf    *bytes.Buffer
   128  	logger log.Logger
   129  }
   130  
   131  func (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  
   142  func (l *colorLogger) putLoggerBuf(cb *loggerBuf) {
   143  	l.bufPool.Put(cb)
   144  }