From 2b5766b4f0cbfd1fe8d0120f058989a2697b5700 Mon Sep 17 00:00:00 2001
From: Simon Sawert <simon@sawert.se>
Date: Sat, 4 Nov 2023 13:16:49 +0100
Subject: [PATCH] Add support for `logr.Marshaler`

---
 example/main.go |  9 +++++++++
 go.sum          |  9 ---------
 logrusr.go      | 51 ++++++++++++++++++++++++++++++-------------------
 logrusr_test.go | 27 ++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 29 deletions(-)

diff --git a/example/main.go b/example/main.go
index 8e98466..f216cd9 100644
--- a/example/main.go
+++ b/example/main.go
@@ -8,6 +8,12 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+type customType string
+
+func (c customType) MarshalLog() interface{} {
+	return "custom marshal message"
+}
+
 func main() {
 	logrusLog := logrus.New()
 	log := logrusr.New(logrusLog)
@@ -75,4 +81,7 @@ func main() {
 	)
 
 	log.Info("custom types", "my_int", myInt(1), "my_other_int", myInt(2))
+
+	myString := customType("original value")
+	log.Info("custom marshaller", "my_string", myString)
 }
diff --git a/go.sum b/go.sum
index 98861d9..c047871 100644
--- a/go.sum
+++ b/go.sum
@@ -5,30 +5,21 @@ github.com/go-logr/logr v1.2.3 h1:2DntVwHkVopvECVRSlL5PSo9eG+cAkDCuckLubN+rq0=
 github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
-github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0=
 github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
-github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 h1:xHms4gcpe1YE7A3yIllJXP16CMAGuqwO2lX1mTyyRRc=
-golang.org/x/sys v0.0.0-20220422013727-9388b58f7150/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A=
 golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
-gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/logrusr.go b/logrusr.go
index c0ded86..3c15199 100644
--- a/logrusr.go
+++ b/logrusr.go
@@ -25,7 +25,10 @@ type FormatFunc func(interface{}) interface{}
 // Option is options to give when construction a logrusr logger.
 type Option func(l *logrusr)
 
-// WithFormatter will set the FormatFunc to use.
+// WithFormatter will set the FormatFunc to use. If you only want to format a
+// specific type that is yours, prefer using the logr.Marshaler interface
+// instead. The `FormatFunc` is better suited for types that are not yours such
+// as external types, maps or slices.
 func WithFormatter(f FormatFunc) Option {
 	return func(l *logrusr) {
 		l.formatter = f
@@ -168,25 +171,33 @@ func listToLogrusFields(formatter FormatFunc, keysAndValues ...interface{}) logr
 	for i := 0; i < len(keysAndValues); i += 2 {
 		k, v := keysAndValues[i], keysAndValues[i+1]
 
-		if s, ok := k.(string); ok {
-			// Try to avoid marshaling known types.
-			switch vVal := v.(type) {
-			case int, int8, int16, int32, int64,
-				uint, uint8, uint16, uint32, uint64,
-				float32, float64, complex64, complex128,
-				string, bool:
-				f[s] = vVal
-
-			case []byte:
-				f[s] = string(vVal)
-
-			default:
-				if formatter != nil {
-					f[s] = formatter(v)
-				} else {
-					j, _ := json.Marshal(vVal)
-					f[s] = string(j)
-				}
+		s, ok := k.(string)
+		if !ok {
+			continue
+		}
+
+		if v, ok := v.(logr.Marshaler); ok {
+			f[s] = v.MarshalLog()
+			continue
+		}
+
+		// Try to avoid marshaling known types.
+		switch vVal := v.(type) {
+		case int, int8, int16, int32, int64,
+			uint, uint8, uint16, uint32, uint64,
+			float32, float64, complex64, complex128,
+			string, bool:
+			f[s] = vVal
+
+		case []byte:
+			f[s] = string(vVal)
+
+		default:
+			if formatter != nil {
+				f[s] = formatter(v)
+			} else {
+				j, _ := json.Marshal(vVal)
+				f[s] = string(j)
 			}
 		}
 	}
diff --git a/logrusr_test.go b/logrusr_test.go
index d52e8fd..6cb6dc5 100644
--- a/logrusr_test.go
+++ b/logrusr_test.go
@@ -15,6 +15,17 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type reverseString string
+
+func (rs reverseString) MarshalLog() interface{} {
+	runes := []rune(rs)
+	for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
+		runes[i], runes[j] = runes[j], runes[i]
+	}
+
+	return string(runes)
+}
+
 func TestLogging(t *testing.T) {
 	t.Parallel()
 
@@ -283,6 +294,22 @@ func TestLogging(t *testing.T) {
 				"caller": `~testing.go:\d+`,
 			},
 		},
+		{
+			description: "custom marshaler used",
+			logFunc: func(log logr.Logger) {
+				log.Info("reversed",
+					"forward", "hello, world",
+					"reversed", reverseString("hello, world"),
+				)
+			},
+			reportCaller: true,
+			assertions: map[string]string{
+				"level":    "info",
+				"msg":      "reversed",
+				"forward":  "hello, world",
+				"reversed": "dlrow ,olleh",
+			},
+		},
 	}
 
 	for _, tc := range cases {