Add support for adding headers

encoding.go

@@ -86,8 +86,8 @@ func (e encoder) writeSource(buf *buffer, pc uintptr, cwd string) {
 		buf.AppendString(frame.File)
 		buf.AppendByte(':')
 		buf.AppendInt(int64(frame.Line))
+		buf.AppendByte(' ')
 	})
-	e.writeColoredString(buf, " > ", e.opts.Theme.AttrKey())
 }
 
 func (e encoder) writeMessage(buf *buffer, level slog.Level, msg string) {
@@ -98,6 +98,26 @@ func (e encoder) writeMessage(buf *buffer, level slog.Level, msg string) {
 	}
 }
 
+func (e encoder) writeHeaders(buf *buffer, headers []slog.Value) bool {
+	wrote := false
+	for _, a := range headers {
+		// todo: this skips empty values, omitting them entire from the header.
+		// alternately, I could print <null> or something, so the number of
+		// headers in each log entry is always fixed...
+		if a.Kind() == slog.KindAny && a.Any() == nil {
+			continue
+		}
+		e.writeValue(buf, a, e.opts.Theme.Source())
+		buf.AppendByte(' ')
+		wrote = true
+	}
+	return wrote
+}
+
+func (e encoder) writeHeaderSeparator(buf *buffer) {
+	e.writeColoredString(buf, "> ", e.opts.Theme.AttrKey())
+}
+
 func (e encoder) writeAttr(buf *buffer, a slog.Attr, group string) {
 	// Elide empty Attrs.
 	if a.Equal(slog.Attr{}) {
@@ -123,38 +143,37 @@ func (e encoder) writeAttr(buf *buffer, a slog.Attr, group string) {
 		buf.AppendString(a.Key)
 		buf.AppendByte('=')
 	})
-	e.writeValue(buf, value)
+	e.writeValue(buf, value, e.opts.Theme.AttrValue())
 }
 
-func (e encoder) writeValue(buf *buffer, value slog.Value) {
-	attrValue := e.opts.Theme.AttrValue()
+func (e encoder) writeValue(buf *buffer, value slog.Value, c ANSIMod) {
 	switch value.Kind() {
 	case slog.KindInt64:
-		e.writeColoredInt(buf, value.Int64(), attrValue)
+		e.writeColoredInt(buf, value.Int64(), c)
 	case slog.KindBool:
-		e.writeColoredBool(buf, value.Bool(), attrValue)
+		e.writeColoredBool(buf, value.Bool(), c)
 	case slog.KindFloat64:
-		e.writeColoredFloat(buf, value.Float64(), attrValue)
+		e.writeColoredFloat(buf, value.Float64(), c)
 	case slog.KindTime:
-		e.writeColoredTime(buf, value.Time(), e.opts.TimeFormat, attrValue)
+		e.writeColoredTime(buf, value.Time(), e.opts.TimeFormat, c)
 	case slog.KindUint64:
-		e.writeColoredUint(buf, value.Uint64(), attrValue)
+		e.writeColoredUint(buf, value.Uint64(), c)
 	case slog.KindDuration:
-		e.writeColoredDuration(buf, value.Duration(), attrValue)
+		e.writeColoredDuration(buf, value.Duration(), c)
 	case slog.KindAny:
 		switch v := value.Any().(type) {
 		case error:
 			e.writeColoredString(buf, v.Error(), e.opts.Theme.AttrValueError())
 			return
 		case fmt.Stringer:
-			e.writeColoredString(buf, v.String(), attrValue)
+			e.writeColoredString(buf, v.String(), c)
 			return
 		}
 		fallthrough
 	case slog.KindString:
 		fallthrough
 	default:
-		e.writeColoredString(buf, value.String(), attrValue)
+		e.writeColoredString(buf, value.String(), c)
 	}
 }
 

handler.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -38,13 +39,19 @@ type HandlerOptions struct {
 
 	// Theme defines the colorized output using ANSI escape sequences
 	Theme Theme
+
+	// Headers are a list of attribute keys.  These attributes will be removed from
+	// the trailing attr list, and the values will be inserted between
+	// the level/source and the message, in the configured order.
+	Headers []string
 }
 
 type Handler struct {
 	opts    HandlerOptions
 	out     io.Writer
 	group   string
 	context buffer
+	headers []slog.Value
 	enc     *encoder
 }
 
@@ -72,6 +79,7 @@ func NewHandler(out io.Writer, opts *HandlerOptions) *Handler {
 		group:   "",
 		context: nil,
 		enc:     &encoder{opts: *opts},
+		headers: make([]slog.Value, len(opts.Headers)),
 	}
 }
 
@@ -82,31 +90,66 @@ func (h *Handler) Enabled(_ context.Context, l slog.Level) bool {
 
 // Handle implements slog.Handler.
 func (h *Handler) Handle(_ context.Context, rec slog.Record) error {
-	buf := bufferPool.Get().(*buffer)
+	headerBuf, trailerBuf := bufferPool.Get().(*buffer), bufferPool.Get().(*buffer)
+
+	h.enc.writeTimestamp(headerBuf, rec.Time)
+	h.enc.writeLevel(headerBuf, rec.Level)
 
-	h.enc.writeTimestamp(buf, rec.Time)
-	h.enc.writeLevel(buf, rec.Level)
+	var writeHeaderSeparator bool
 	if h.opts.AddSource && rec.PC > 0 {
-		h.enc.writeSource(buf, rec.PC, cwd)
+		h.enc.writeSource(headerBuf, rec.PC, cwd)
+		writeHeaderSeparator = true
 	}
-	h.enc.writeMessage(buf, rec.Level, rec.Message)
-	buf.copy(&h.context)
+
+	h.enc.writeMessage(trailerBuf, rec.Level, rec.Message)
+
+	trailerBuf.copy(&h.context)
+
+	headers := h.headers
+	headersChanged := false
 	rec.Attrs(func(a slog.Attr) bool {
-		h.enc.writeAttr(buf, a, h.group)
+		idx := slices.IndexFunc(h.opts.Headers, func(s string) bool { return s == a.Key })
+		if idx >= 0 {
+			if !headersChanged {
+				headersChanged = true
+				// todo: this makes one allocation, but only if the headers weren't already
+				// satisfied by prior WithAttrs().  Could use a pool of *[]slog.Value, but
+				// I'm not sure it's worth it.
+				headers = make([]slog.Value, len(h.opts.Headers))
+				copy(headers, h.headers)
+			}
+			headers[idx] = a.Value
+		} else {
+			h.enc.writeAttr(trailerBuf, a, h.group)
+		}
 		return true
 	})
-	h.enc.NewLine(buf)
-	if _, err := buf.WriteTo(h.out); err != nil {
-		buf.Reset()
-		bufferPool.Put(buf)
+	h.enc.NewLine(trailerBuf)
+
+	if len(headers) > 0 {
+		if h.enc.writeHeaders(headerBuf, headers) {
+			writeHeaderSeparator = true
+		}
+	}
+
+	if writeHeaderSeparator {
+		h.enc.writeHeaderSeparator(headerBuf)
+	}
+
+	if _, err := headerBuf.WriteTo(h.out); err != nil {
+		return err
+	}
+	if _, err := trailerBuf.WriteTo(h.out); err != nil {
 		return err
 	}
-	bufferPool.Put(buf)
+	bufferPool.Put(headerBuf)
+	bufferPool.Put(trailerBuf)
 	return nil
 }
 
 // WithAttrs implements slog.Handler.
 func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	headers := h.extractHeaders(attrs)
 	newCtx := h.context
 	for _, a := range attrs {
 		h.enc.writeAttr(&newCtx, a, h.group)
@@ -118,6 +161,7 @@ func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler {
 		group:   h.group,
 		context: newCtx,
 		enc:     h.enc,
+		headers: headers,
 	}
 }
 
@@ -133,5 +177,28 @@ func (h *Handler) WithGroup(name string) slog.Handler {
 		group:   name,
 		context: h.context,
 		enc:     h.enc,
+		headers: h.headers,
+	}
+}
+
+// extractHeaders scans the attributes for keys specified in Headers.
+// If found, their values are saved in a new list.
+// The original attribute list will be modified to remove the extracted attributes.
+func (h *Handler) extractHeaders(attrs []slog.Attr) (headers []slog.Value) {
+	changed := false
+	headers = h.headers
+	for i, attr := range attrs {
+		idx := slices.IndexFunc(h.opts.Headers, func(s string) bool { return s == attr.Key })
+		if idx >= 0 {
+			if !changed {
+				// make a copy of prefixes:
+				headers = make([]slog.Value, len(h.headers))
+				copy(headers, h.headers)
+			}
+			headers[idx] = attr.Value
+			attrs[i] = slog.Attr{} // remove the prefix attribute
+			changed = true
+		}
 	}
+	return
 }

handler_test.go

@@ -271,6 +271,152 @@ func TestHandler_Source(t *testing.T) {
 	AssertEqual(t, fmt.Sprintf("%s INF foobar\n", now.Format(time.DateTime)), buf.String())
 }
 
+func TestHandler_Headers(t *testing.T) {
+	pc, file, line, _ := runtime.Caller(0)
+	cwd, _ := os.Getwd()
+	file, _ = filepath.Rel(cwd, file)
+	sourceField := fmt.Sprintf("%s:%d", file, line)
+
+	tests := []struct {
+		name       string
+		opts       HandlerOptions
+		attrs      []slog.Attr
+		withAttrs  []slog.Attr
+		withGroups []string
+		want       string
+	}{
+		{
+			name:  "no headers",
+			attrs: []slog.Attr{slog.String("foo", "bar")},
+			want:  "INF with headers foo=bar\n",
+		},
+		{
+			name: "one header",
+			opts: HandlerOptions{Headers: []string{"foo"}},
+			attrs: []slog.Attr{
+				slog.String("foo", "bar"),
+				slog.String("bar", "baz"),
+			},
+			want: "INF bar > with headers bar=baz\n",
+		},
+		{
+			name: "two headers",
+			opts: HandlerOptions{Headers: []string{"foo", "bar"}},
+			attrs: []slog.Attr{
+				slog.String("foo", "bar"),
+				slog.String("bar", "baz"),
+			},
+			want: "INF bar baz > with headers\n",
+		},
+		{
+			name:  "missing headers",
+			opts:  HandlerOptions{Headers: []string{"foo", "bar"}},
+			attrs: []slog.Attr{slog.String("bar", "baz"), slog.String("baz", "foo")},
+			want:  "INF baz > with headers baz=foo\n",
+		},
+		{
+			name: "missing all headers",
+			opts: HandlerOptions{Headers: []string{"foo", "bar"}},
+			want: "INF with headers\n",
+		},
+		{
+			name: "header and source",
+			opts: HandlerOptions{Headers: []string{"foo"}, AddSource: true},
+			attrs: []slog.Attr{
+				slog.String("foo", "bar"),
+				slog.String("bar", "baz"),
+			},
+			want: "INF " + sourceField + " bar > with headers bar=baz\n",
+		},
+		{
+			name: "withattrs",
+			opts: HandlerOptions{Headers: []string{"foo"}},
+			attrs: []slog.Attr{
+
+				slog.String("bar", "baz"),
+			},
+			withAttrs: []slog.Attr{
+				slog.String("foo", "bar"),
+			},
+			want: "INF bar > with headers bar=baz\n",
+		},
+		{
+			name: "withgroup",
+			opts: HandlerOptions{Headers: []string{"foo", "bar"}},
+			attrs: []slog.Attr{
+				slog.String("bar", "baz"),
+				slog.String("baz", "foo"),
+			},
+			withGroups: []string{"group"},
+			withAttrs: []slog.Attr{
+				slog.String("foo", "bar"),
+			},
+			want: "INF bar baz > with headers group.baz=foo\n",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			buf := bytes.Buffer{}
+
+			opts := &test.opts
+			opts.NoColor = true
+			var h slog.Handler = NewHandler(&buf, &test.opts)
+			if test.withAttrs != nil {
+				h = h.WithAttrs(test.withAttrs)
+			}
+			for _, g := range test.withGroups {
+				h = h.WithGroup(g)
+			}
+
+			rec := slog.NewRecord(time.Time{}, slog.LevelInfo, "with headers", pc)
+
+			rec.AddAttrs(test.attrs...)
+
+			AssertNoError(t, h.Handle(context.Background(), rec))
+			AssertEqual(t, test.want, buf.String())
+		})
+	}
+
+	t.Run("withAttrs state keeping", func(t *testing.T) {
+		// test to make sure the way that WithAttrs() copies the cached headers doesn't leak
+		// headers back to the parent handler or to subsequent Handle() calls (i.e. ensure that
+		// the headers slice is copied at the right times).
+
+		buf := bytes.Buffer{}
+		h := NewHandler(&buf, &HandlerOptions{
+			Headers:    []string{"foo", "bar"},
+			TimeFormat: "0",
+			NoColor:    true,
+		})
+
+		assertLog := func(t *testing.T, handler slog.Handler, want string, attrs ...slog.Attr) {
+			buf.Reset()
+			rec := slog.NewRecord(time.Time{}, slog.LevelInfo, "with headers", pc)
+
+			rec.AddAttrs(attrs...)
+
+			AssertNoError(t, handler.Handle(context.Background(), rec))
+			AssertEqual(t, want, buf.String())
+		}
+
+		assertLog(t, h, "INF bar > with headers\n", slog.String("foo", "bar"))
+
+		h2 := h.WithAttrs([]slog.Attr{slog.String("foo", "baz")})
+		assertLog(t, h2, "INF baz > with headers\n")
+
+		h3 := h2.WithAttrs([]slog.Attr{slog.String("foo", "buz")})
+		assertLog(t, h3, "INF buz > with headers\n")
+		// creating h3 should not have affected h2
+		assertLog(t, h2, "INF baz > with headers\n")
+
+		// overriding attrs shouldn't affect the handler
+		assertLog(t, h2, "INF biz > with headers\n", slog.String("foo", "biz"))
+		assertLog(t, h2, "INF baz > with headers\n")
+
+	})
+}
+
 func TestHandler_Err(t *testing.T) {
 	w := writerFunc(func(b []byte) (int, error) { return 0, errors.New("nope") })
 	h := NewHandler(w, &HandlerOptions{NoColor: true})