Add FromContext and WithContext (#11)

* Add `FromContext` and `WithContext`

* Add tests for the missing logger case

Author: FollowTheProcess

log.go

@@ -10,8 +10,10 @@ package log
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
+	"os"
 	"slices"
 	"strconv"
 	"sync"
@@ -34,8 +36,15 @@ const (
 	errorStyle     = hue.Red | hue.Bold
 )
 
+// missingValue is the placeholder text for a missing value in a log line's key value pair.
 const missingValue = "<MISSING>"
 
+// ctxKey is the unexported type used for context key so this key never collides with another.
+type ctxKey struct{}
+
+// contextKey is the actual key used to store and retrieve a Logger from a Context.
+var contextKey = ctxKey{}
+
 // Logger is a command line logger. It is safe to use across concurrently
 // executing goroutines.
 type Logger struct {
@@ -58,7 +67,7 @@ func New(w io.Writer, options ...Option) *Logger {
 		w:          w,
 		level:      LevelInfo,
 		timeFormat: time.RFC3339,
-		timeFunc:   now,
+		timeFunc:   func() time.Time { return time.Now().UTC() },
 		mu:         &sync.Mutex{},
 	}
 
@@ -71,6 +80,25 @@ func New(w io.Writer, options ...Option) *Logger {
 	return logger
 }
 
+// WithContext stores the given logger in a [context.Context].
+//
+// The logger may be retrieved from the context with [FromContext].
+func WithContext(ctx context.Context, logger *Logger) context.Context {
+	return context.WithValue(ctx, contextKey, logger)
+}
+
+// FromContext returns the [Logger] from a [context.Context].
+//
+// If the context does not contain a logger, a default logger is returned.
+func FromContext(ctx context.Context) *Logger {
+	logger, ok := ctx.Value(contextKey).(*Logger)
+	if !ok || logger == nil {
+		return New(os.Stderr)
+	}
+
+	return logger
+}
+
 // With returns a new [Logger] with the given persistent key value pairs.
 //
 // The returned logger is otherwise an exact clone of the caller.
@@ -220,11 +248,6 @@ func putBuffer(buf *bytes.Buffer) {
 	bufPool.Put(buf)
 }
 
-// now returns the current time with UTC.
-func now() time.Time {
-	return time.Now().UTC()
-}
-
 // needsQuotes returns whether s should be displayed as "s".
 func needsQuotes(s string) bool {
 	for _, char := range s {

log_test.go

@@ -214,6 +214,46 @@ func TestRace(t *testing.T) {
 	test.Equal(t, len(lines), n*2, test.Context("expected %d log lines", n*2))
 }
 
+func TestContext(t *testing.T) {
+	t.Run("present", func(t *testing.T) {
+		buf := &bytes.Buffer{}
+
+		// Constantly return the same time
+		fixedTime := func() time.Time {
+			fixed, err := time.Parse(time.RFC3339, "2025-04-01T13:34:03Z")
+			test.Ok(t, err)
+
+			return fixed
+		}
+
+		// Configure it a bit so we know we're getting the right one
+		logger := log.New(buf, log.TimeFunc(fixedTime), log.TimeFormat(time.Kitchen))
+
+		logger.Info("Before")
+
+		ctx := t.Context()
+
+		ctx = log.WithContext(ctx, logger)
+
+		after := log.FromContext(ctx)
+
+		after.Info("After")
+
+		got := buf.String()
+
+		test.Diff(t, got, "1:34PM INFO: Before\n1:34PM INFO: After\n")
+	})
+
+	t.Run("missing", func(t *testing.T) {
+		_, stderr := test.CaptureOutput(t, func() error {
+			log.FromContext(t.Context()).Info("FromContext")
+			return nil
+		})
+
+		test.True(t, strings.Contains(stderr, "FromContext"))
+	})
+}
+
 func BenchmarkLogger(b *testing.B) {
 	hue.Enabled(true) // Force colour