Extend WithLazy() method to SugaredLogger

sugar.go

@@ -115,6 +115,21 @@ func (s *SugaredLogger) With(args ...interface{}) *SugaredLogger {
 	return &SugaredLogger{base: s.base.With(s.sweetenFields(args)...)}
 }
 
+// WithLazy adds a variadic number of fields to the logging context lazily.
+// The fields are evaluated only if the logger is further chained with [With]
+// or is written to with any of the log level methods.
+// Until that occurs, the logger may retain references to objects inside the fields,
+// and logging will reflect the state of an object at the time of logging,
+// not the time of WithLazy().
+//
+// Similar to [With], fields added to the child don't affect the parent,
+// and vice versa. Also, the keys in key-value pairs should be strings. In development,
+// passing a non-string key panics, while in production it logs an error and skips the pair.
+// Passing an orphaned key has the same behavior.
+func (s *SugaredLogger) WithLazy(args ...interface{}) *SugaredLogger {
+	return &SugaredLogger{base: s.base.WithLazy(s.sweetenFields(args)...)}
+}
+
 // Level reports the minimum enabled level for this logger.
 //
 // For NopLoggers, this is [zapcore.InvalidLevel].

sugar_test.go

@@ -22,6 +22,8 @@ package zap
 
 import (
 	"errors"
+	"fmt"
+	"strconv"
 	"testing"
 
 	"go.uber.org/zap/internal/exit"
@@ -54,6 +56,9 @@ func TestSugarWith(t *testing.T) {
 		}
 	}
 
+	type withAny func(*SugaredLogger, ...interface{}) *SugaredLogger
+	withMethods := []withAny{(*SugaredLogger).With, (*SugaredLogger).WithLazy}
+
 	tests := []struct {
 		desc     string
 		args     []interface{}
@@ -141,16 +146,94 @@ func TestSugarWith(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		withSugar(t, DebugLevel, nil, func(logger *SugaredLogger, logs *observer.ObservedLogs) {
-			logger.With(tt.args...).Info("")
-			output := logs.AllUntimed()
-			if len(tt.errLogs) > 0 {
-				for i := range tt.errLogs {
-					assert.Equal(t, tt.errLogs[i], output[i], "Unexpected error log at position %d for scenario %s.", i, tt.desc)
+		for _, withMethod := range withMethods {
+			withSugar(t, DebugLevel, nil, func(logger *SugaredLogger, logs *observer.ObservedLogs) {
+				withMethod(logger, tt.args...).Info("")
+				output := logs.AllUntimed()
+				if len(tt.errLogs) > 0 {
+					for i := range tt.errLogs {
+						assert.Equal(t, tt.errLogs[i], output[i], "Unexpected error log at position %d for scenario %s.", i, tt.desc)
+					}
+				}
+				assert.Equal(t, len(tt.errLogs)+1, len(output), "Expected only one non-error message to be logged in scenario %s.", tt.desc)
+				assert.Equal(t, tt.expected, output[len(tt.errLogs)].Context, "Unexpected message context in scenario %s.", tt.desc)
+			})
+		}
+	}
+}
+
+func TestSugarWithCaptures(t *testing.T) {
+	type withAny func(*SugaredLogger, ...interface{}) *SugaredLogger
+
+	tests := []struct {
+		name        string
+		withMethods []withAny
+		wantJSON    []string
+	}{
+		{
+			name:        "with captures arguments at time of With",
+			withMethods: []withAny{(*SugaredLogger).With},
+			wantJSON: []string{
+				`{
+					"m": "hello 0",
+					"a0": [0],
+					"b0": [1]
+				}`,
+				`{
+					"m": "world 0",
+					"a0": [0],
+					"c0": [2]
+				}`,
+			},
+		},
+		{
+			name:        "lazy with captures arguments at time of Logging",
+			withMethods: []withAny{(*SugaredLogger).WithLazy},
+			wantJSON: []string{
+				`{
+					"m": "hello 0",
+					"a0": [1],
+					"b0": [1]
+				}`,
+				`{
+					"m": "world 0",
+					"a0": [1],
+					"c0": [2]
+				}`,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			enc := zapcore.NewJSONEncoder(zapcore.EncoderConfig{
+				MessageKey: "m",
+			})
+
+			var bs ztest.Buffer
+			logger := New(zapcore.NewCore(enc, &bs, DebugLevel)).Sugar()
+
+			for i, withMethod := range tt.withMethods {
+				iStr := strconv.Itoa(i)
+				x := 10 * i
+				arr := zapcore.ArrayMarshalerFunc(func(enc zapcore.ArrayEncoder) error {
+					enc.AppendInt(x)
+					return nil
+				})
+
+				logger = withMethod(logger, Array("a"+iStr, arr))
+				x++
+				logger.Infow(fmt.Sprintf("hello %d", i), Array("b"+iStr, arr))
+				x++
+				logger = withMethod(logger, Array("c"+iStr, arr))
+				logger.Infow(fmt.Sprintf("world %d", i))
+			}
+
+			if lines := bs.Lines(); assert.Len(t, lines, len(tt.wantJSON)) {
+				for i, want := range tt.wantJSON {
+					assert.JSONEq(t, want, lines[i], "Unexpected output from the %d'th log.", i)
 				}
 			}
-			assert.Equal(t, len(tt.errLogs)+1, len(output), "Expected only one non-error message to be logged in scenario %s.", tt.desc)
-			assert.Equal(t, tt.expected, output[len(tt.errLogs)].Context, "Unexpected message context in scenario %s.", tt.desc)
 		})
 	}
 }