colorlog.go

colorlog.go

// Copyright (c) Gabriel de Quadros Ligneul
// SPDX-License-Identifier: Apache-2.0 (see LICENSE)

// This module implements a colored log handler for log/slog.
//
// For more details on a custom handler implementation, check:
// - https://github.com/golang/example/tree/master/slog-handler-guide
package colorlog

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"runtime"
	"slices"
	"sync"
)

// ANSI color code
type Color string

const (
	Black         Color = "30"
	Red           Color = "31"
	Green         Color = "32"
	Yellow        Color = "33"
	Blue          Color = "34"
	Magenta       Color = "35"
	Cyan          Color = "36"
	White         Color = "37"
	BrightBlack   Color = "90"
	BrightRed     Color = "91"
	BrightGreen   Color = "92"
	BrightYellow  Color = "93"
	BrightBlue    Color = "94"
	BrightMagenta Color = "95"
	BrightCyan    Color = "96"
	BrightWhite   Color = "97"
)

const (
	DefaultTimeLayout = "[15:04:05.000]"
)

type Options struct {
	// AddSource causes the handler to compute the source code position
	// of the log statement and add a SourceKey attribute to the output.
	AddSource bool

	// Level reports the minimum record level that will be logged.
	// The handler discards records with lower levels.
	// If Level is nil, the handler assumes LevelInfo.
	// The handler calls Level.Level for each record processed;
	// to adjust the minimum level dynamically, use a LevelVar.
	Level slog.Leveler

	// EnableColor defines wheter the handler prints logs with color.
	EnableColor bool

	// TimeLayout is the layout used to print the log time.
	TimeLayout string
}

type groupedAttr struct {
	groups []string
	attr   slog.Attr
}

type Handler struct {
	opts   Options
	out    io.Writer
	mu     *sync.Mutex   // must be a pointer because we copy it in WithGroup and WithAttrs
	groups []string      // group names from WithGroup
	attrs  []groupedAttr // attrs from WithAttrs
}

func NewHandler(out io.Writer, opts *Options) *Handler {
	h := &Handler{
		out:    out,
		mu:     new(sync.Mutex),
		groups: nil,
		attrs:  nil,
	}
	if opts != nil {
		h.opts = *opts
	}
	if h.opts.Level == nil {
		h.opts.Level = slog.LevelInfo
	}
	if h.opts.TimeLayout == "" {
		h.opts.TimeLayout = DefaultTimeLayout
	}
	return h
}

func (h *Handler) Enabled(ctx context.Context, level slog.Level) bool {
	return level >= h.opts.Level.Level()
}

func (h *Handler) WithGroup(name string) slog.Handler {
	if name == "" {
		return h
	}
	h2 := *h
	groups := slices.Clone(h.groups) // clone to avoid aliasing
	h2.groups = append(groups, name)
	return &h2
}

func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return h
	}
	h2 := *h
	h2.attrs = slices.Clone(h.attrs) // clone to avoid aliasing
	for i := range attrs {
		h2.attrs = append(h2.attrs, groupedAttr{groups: h.groups, attr: attrs[i]})
	}
	return &h2
}

func (h *Handler) Handle(ctx context.Context, r slog.Record) error {
	buf := bytes.NewBuffer(make([]byte, 0, 1024))

	// timestamp
	h.colorized(buf, BrightBlack, r.Time.Format(h.opts.TimeLayout))
	buf.WriteRune(' ')

	// level
	var levelColor Color
	if r.Level <= slog.LevelDebug {
		levelColor = Cyan
	} else if r.Level == slog.LevelInfo {
		levelColor = Green
	} else if r.Level == slog.LevelWarn {
		levelColor = Yellow
	} else { // error or above
		levelColor = Red
	}
	h.colorized(buf, levelColor, r.Level.String())
	buf.WriteRune(' ')

	// message
	buf.WriteString(r.Message)

	// attributes
	for i := range h.attrs {
		buf.WriteString(" ")
		h.writeAttr(buf, Blue, h.attrs[i].groups, h.attrs[i].attr)
	}
	r.Attrs(func(attr slog.Attr) bool {
		buf.WriteString(" ")
		h.writeAttr(buf, Blue, h.groups, attr)
		return true
	})
	buf.WriteRune('\n')

	// debug info
	if h.opts.AddSource && r.PC != 0 {
		fs := runtime.CallersFrames([]uintptr{r.PC})
		f, _ := fs.Next()
		groups := []string{"source"}
		buf.WriteString("    ")
		h.writeAttr(buf, Magenta, groups, slog.String("function", f.Function))
		buf.WriteString("\n    ")
		h.writeAttr(buf, Magenta, groups, slog.String("file", f.File))
		buf.WriteString("\n    ")
		h.writeAttr(buf, Magenta, groups, slog.Int("line", f.Line))
		buf.WriteRune('\n')
	}

	// write to output
	h.mu.Lock()
	defer h.mu.Unlock()
	_, err := h.out.Write(buf.Bytes())
	return err
}

// Auxiliary functions -----------------------------------------------------------------------------

// Write the attribute to the buffer.
func (h *Handler) writeAttr(buf *bytes.Buffer, color Color, groups []string, attr slog.Attr) {
	attr.Value = attr.Value.Resolve()

	// ignore empty attributes
	if attr.Equal(slog.Attr{}) {
		return
	}

	// handle groups recursively
	if attr.Value.Kind() == slog.KindGroup {
		attrs := attr.Value.Group()
		// ignore empty groups
		if len(attrs) == 0 {
			return
		}
		if attr.Key != "" {
			groups = slices.Clone(h.groups) // clone to avoid aliasing
			groups = append(groups, attr.Key)
		}
		for i := range attrs {
			h.writeAttr(buf, color, groups, attrs[i])
		}
		return
	}

	h.startColor(buf, color)

	// write groups
	for i := range groups {
		buf.WriteString(groups[i])
		buf.WriteRune('.')
	}

	// write key
	if attr.Key != "" {
		buf.WriteString(attr.Key)
	} else {
		buf.WriteString("NOKEY")
	}

	// write value
	buf.WriteRune('=')
	buf.WriteString(attr.Value.String())

	h.endColor(buf)
}

// Wrap the string with the appropriate color code.
func (h *Handler) colorized(buf *bytes.Buffer, color Color, v string) {
	h.startColor(buf, color)
	buf.WriteString(v)
	h.endColor(buf)
}

func (h *Handler) startColor(buf *bytes.Buffer, color Color) {
	if h.opts.EnableColor {
		buf.WriteString("\033[")
		buf.WriteString(string(color))
		buf.WriteString("m")
	}
}

func (h *Handler) endColor(buf *bytes.Buffer) {
	if h.opts.EnableColor {
		buf.WriteString("\033[0m")
	}
}