Update custom marshaler and unmarshaler to accept context (#745)

* feat: Update custom marshaler and unmarshaler to accept context

* Add tests for custom marshaler and unmarshaler with context

Author: pgeiem

decode.go

@@ -30,7 +30,7 @@ type Decoder struct {
 	referenceReaders     []io.Reader
 	anchorNodeMap        map[string]ast.Node
 	anchorValueMap       map[string]reflect.Value
-	customUnmarshalerMap map[reflect.Type]func(interface{}, []byte) error
+	customUnmarshalerMap map[reflect.Type]func(context.Context, interface{}, []byte) error
 	commentMaps          []CommentMap
 	toCommentMap         CommentMap
 	opts                 []DecodeOption
@@ -54,7 +54,7 @@ func NewDecoder(r io.Reader, opts ...DecodeOption) *Decoder {
 		reader:               r,
 		anchorNodeMap:        map[string]ast.Node{},
 		anchorValueMap:       map[string]reflect.Value{},
-		customUnmarshalerMap: map[reflect.Type]func(interface{}, []byte) error{},
+		customUnmarshalerMap: map[reflect.Type]func(context.Context, interface{}, []byte) error{},
 		opts:                 opts,
 		referenceReaders:     []io.Reader{},
 		referenceFiles:       []string{},
@@ -722,7 +722,7 @@ func (d *Decoder) existsTypeInCustomUnmarshalerMap(t reflect.Type) bool {
 	return false
 }
 
-func (d *Decoder) unmarshalerFromCustomUnmarshalerMap(t reflect.Type) (func(interface{}, []byte) error, bool) {
+func (d *Decoder) unmarshalerFromCustomUnmarshalerMap(t reflect.Type) (func(context.Context, interface{}, []byte) error, bool) {
 	if unmarshaler, exists := d.customUnmarshalerMap[t]; exists {
 		return unmarshaler, exists
 	}
@@ -765,7 +765,7 @@ func (d *Decoder) decodeByUnmarshaler(ctx context.Context, dst reflect.Value, sr
 		if err != nil {
 			return err
 		}
-		if err := unmarshaler(ptrValue.Interface(), b); err != nil {
+		if err := unmarshaler(ctx, ptrValue.Interface(), b); err != nil {
 			return err
 		}
 		return nil

decode_test.go

@@ -2173,6 +2173,29 @@ func TestDecoder_CustomUnmarshaler(t *testing.T) {
 			t.Fatalf("failed to switch to custom unmarshaler. got: %q", v.Foo)
 		}
 	})
+	t.Run("override bytes type with context", func(t *testing.T) {
+		type T struct {
+			Foo []byte `yaml:"foo"`
+		}
+		src := []byte(`foo: "bar"`)
+		var v T
+		ctx := context.WithValue(context.Background(), "plop", uint(42))
+		if err := yaml.UnmarshalContext(ctx, src, &v, yaml.CustomUnmarshalerContext[[]byte](func(ctx context.Context, dst *[]byte, b []byte) error {
+			if !bytes.Equal(b, []byte(`"bar"`)) {
+				t.Fatalf("failed to get target buffer: %q", b)
+			}
+			if ctx.Value("plop") != uint(42) {
+				t.Fatalf("context value is not correct")
+			}
+			*dst = []byte("bazbaz")
+			return nil
+		})); err != nil {
+			t.Fatal(err)
+		}
+		if !bytes.Equal(v.Foo, []byte("bazbaz")) {
+			t.Fatalf("failed to switch to custom unmarshaler. got: %q", v.Foo)
+		}
+	})
 }
 
 type unmarshalContext struct {

encode.go

@@ -37,7 +37,7 @@ type Encoder struct {
 	anchorRefToName            map[uintptr]string
 	anchorNameMap              map[string]struct{}
 	anchorCallback             func(*ast.AnchorNode, interface{}) error
-	customMarshalerMap         map[reflect.Type]func(interface{}) ([]byte, error)
+	customMarshalerMap         map[reflect.Type]func(context.Context, interface{}) ([]byte, error)
 	omitZero                   bool
 	omitEmpty                  bool
 	autoInt                    bool
@@ -59,7 +59,7 @@ func NewEncoder(w io.Writer, opts ...EncodeOption) *Encoder {
 	return &Encoder{
 		writer:             w,
 		opts:               opts,
-		customMarshalerMap: map[reflect.Type]func(interface{}) ([]byte, error){},
+		customMarshalerMap: map[reflect.Type]func(context.Context, interface{}) ([]byte, error){},
 		line:               1,
 		column:             1,
 		offset:             0,
@@ -301,7 +301,7 @@ func (e *Encoder) existsTypeInCustomMarshalerMap(t reflect.Type) bool {
 	return false
 }
 
-func (e *Encoder) marshalerFromCustomMarshalerMap(t reflect.Type) (func(interface{}) ([]byte, error), bool) {
+func (e *Encoder) marshalerFromCustomMarshalerMap(t reflect.Type) (func(context.Context, interface{}) ([]byte, error), bool) {
 	if marshaler, exists := e.customMarshalerMap[t]; exists {
 		return marshaler, exists
 	}
@@ -347,7 +347,7 @@ func (e *Encoder) encodeByMarshaler(ctx context.Context, v reflect.Value, column
 	iface := v.Interface()
 
 	if marshaler, exists := e.marshalerFromCustomMarshalerMap(v.Type()); exists {
-		doc, err := marshaler(iface)
+		doc, err := marshaler(ctx, iface)
 		if err != nil {
 			return nil, err
 		}

encode_test.go

@@ -1492,6 +1492,27 @@ func TestEncoder_CustomMarshaler(t *testing.T) {
 			t.Fatalf("failed to switch to custom marshaler. got: %q", b)
 		}
 	})
+	t.Run("override bytes type with context", func(t *testing.T) {
+		type T struct {
+			Foo []byte `yaml:"foo"`
+		}
+		ctx := context.WithValue(context.Background(), "plop", uint(42))
+		b, err := yaml.MarshalContext(ctx, &T{Foo: []byte("bar")}, yaml.CustomMarshalerContext[[]byte](func(ctx context.Context, v []byte) ([]byte, error) {
+			if !bytes.Equal(v, []byte("bar")) {
+				t.Fatalf("failed to get src buffer: %q", v)
+			}
+			if ctx.Value("plop") != uint(42) {
+				t.Fatalf("context value is not correct")
+			}
+			return []byte(`override`), nil
+		}))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !bytes.Equal(b, []byte("foo: override\n")) {
+			t.Fatalf("failed to switch to custom marshaler. got: %q", b)
+		}
+	})
 }
 
 func TestEncoder_AutoInt(t *testing.T) {

option.go

@@ -1,6 +1,7 @@
 package yaml
 
 import (
+	"context"
 	"io"
 	"reflect"
 
@@ -101,13 +102,25 @@ func UseJSONUnmarshaler() DecodeOption {
 func CustomUnmarshaler[T any](unmarshaler func(*T, []byte) error) DecodeOption {
 	return func(d *Decoder) error {
 		var typ *T
-		d.customUnmarshalerMap[reflect.TypeOf(typ)] = func(v interface{}, b []byte) error {
+		d.customUnmarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}, b []byte) error {
 			return unmarshaler(v.(*T), b)
 		}
 		return nil
 	}
 }
 
+// CustomUnmarshalerContext overrides any decoding process for the type specified in generics.
+// Similar to CustomUnmarshaler, but allows passing a context to the unmarshaler function.
+func CustomUnmarshalerContext[T any](unmarshaler func(context.Context, *T, []byte) error) DecodeOption {
+	return func(d *Decoder) error {
+		var typ *T
+		d.customUnmarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}, b []byte) error {
+			return unmarshaler(ctx, v.(*T), b)
+		}
+		return nil
+	}
+}
+
 // EncodeOption functional option type for Encoder
 type EncodeOption func(e *Encoder) error
 
@@ -199,13 +212,25 @@ func UseJSONMarshaler() EncodeOption {
 func CustomMarshaler[T any](marshaler func(T) ([]byte, error)) EncodeOption {
 	return func(e *Encoder) error {
 		var typ T
-		e.customMarshalerMap[reflect.TypeOf(typ)] = func(v interface{}) ([]byte, error) {
+		e.customMarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}) ([]byte, error) {
 			return marshaler(v.(T))
 		}
 		return nil
 	}
 }
 
+// CustomMarshalerContext overrides any encoding process for the type specified in generics.
+// Similar to CustomMarshaler, but allows passing a context to the marshaler function.
+func CustomMarshalerContext[T any](marshaler func(context.Context, T) ([]byte, error)) EncodeOption {
+	return func(e *Encoder) error {
+		var typ T
+		e.customMarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}) ([]byte, error) {
+			return marshaler(ctx, v.(T))
+		}
+		return nil
+	}
+}
+
 // AutoInt automatically converts floating-point numbers to integers when the fractional part is zero.
 // For example, a value of 1.0 will be encoded as 1.
 func AutoInt() EncodeOption {

testdata/yaml_test.go

@@ -2,6 +2,7 @@ package yaml_test
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"strings"
 	"testing"
@@ -1084,6 +1085,26 @@ func TestRegisterCustomMarshaler(t *testing.T) {
 	}
 }
 
+func TestRegisterCustomMarshalerContext(t *testing.T) {
+	type T struct {
+		Foo []byte `yaml:"foo"`
+	}
+	yaml.RegisterCustomMarshalerContext[T](func(ctx context.Context, _ T) ([]byte, error) {
+		if ctx.Value("plop") != uint(42) {
+			t.Fatalf("context value is not correct")
+		}
+		return []byte(`"override"`), nil
+	})
+	ctx := context.WithValue(context.Background(), "plop", uint(42))
+	b, err := yaml.MarshalContext(ctx, &T{Foo: []byte("bar")})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(b, []byte("\"override\"\n")) {
+		t.Fatalf("failed to register custom marshaler. got: %q", b)
+	}
+}
+
 func TestRegisterCustomUnmarshaler(t *testing.T) {
 	type T struct {
 		Foo []byte `yaml:"foo"`
@@ -1100,3 +1121,24 @@ func TestRegisterCustomUnmarshaler(t *testing.T) {
 		t.Fatalf("failed to decode. got %q", v.Foo)
 	}
 }
+
+func TestRegisterCustomUnmarshalerContext(t *testing.T) {
+	type T struct {
+		Foo []byte `yaml:"foo"`
+	}
+	yaml.RegisterCustomUnmarshalerContext[T](func(ctx context.Context, v *T, _ []byte) error {
+		if ctx.Value("plop") != uint(42) {
+			t.Fatalf("context value is not correct")
+		}
+		v.Foo = []byte("override")
+		return nil
+	})
+	var v T
+	ctx := context.WithValue(context.Background(), "plop", uint(42))
+	if err := yaml.UnmarshalContext(ctx, []byte(`"foo": "bar"`), &v); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(v.Foo, []byte("override")) {
+		t.Fatalf("failed to decode. got %q", v.Foo)
+	}
+}

yaml.go

@@ -266,8 +266,8 @@ func JSONToYAML(bytes []byte) ([]byte, error) {
 var (
 	globalCustomMarshalerMu    sync.Mutex
 	globalCustomUnmarshalerMu  sync.Mutex
-	globalCustomMarshalerMap   = map[reflect.Type]func(interface{}) ([]byte, error){}
-	globalCustomUnmarshalerMap = map[reflect.Type]func(interface{}, []byte) error{}
+	globalCustomMarshalerMap   = map[reflect.Type]func(context.Context, interface{}) ([]byte, error){}
+	globalCustomUnmarshalerMap = map[reflect.Type]func(context.Context, interface{}, []byte) error{}
 )
 
 // RegisterCustomMarshaler overrides any encoding process for the type specified in generics.
@@ -281,11 +281,23 @@ func RegisterCustomMarshaler[T any](marshaler func(T) ([]byte, error)) {
 	defer globalCustomMarshalerMu.Unlock()
 
 	var typ T
-	globalCustomMarshalerMap[reflect.TypeOf(typ)] = func(v interface{}) ([]byte, error) {
+	globalCustomMarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}) ([]byte, error) {
 		return marshaler(v.(T))
 	}
 }
 
+// RegisterCustomMarshalerContext overrides any encoding process for the type specified in generics.
+// Similar to RegisterCustomMarshalerContext, but allows passing a context to the unmarshaler function.
+func RegisterCustomMarshalerContext[T any](marshaler func(context.Context, T) ([]byte, error)) {
+	globalCustomMarshalerMu.Lock()
+	defer globalCustomMarshalerMu.Unlock()
+
+	var typ T
+	globalCustomMarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}) ([]byte, error) {
+		return marshaler(ctx, v.(T))
+	}
+}
+
 // RegisterCustomUnmarshaler overrides any decoding process for the type specified in generics.
 // If you want to switch the behavior for each decoder, use `CustomUnmarshaler` defined as DecodeOption.
 //
@@ -296,7 +308,19 @@ func RegisterCustomUnmarshaler[T any](unmarshaler func(*T, []byte) error) {
 	defer globalCustomUnmarshalerMu.Unlock()
 
 	var typ *T
-	globalCustomUnmarshalerMap[reflect.TypeOf(typ)] = func(v interface{}, b []byte) error {
+	globalCustomUnmarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}, b []byte) error {
 		return unmarshaler(v.(*T), b)
 	}
 }
+
+// RegisterCustomUnmarshalerContext overrides any decoding process for the type specified in generics.
+// Similar to RegisterCustomUnmarshalerContext, but allows passing a context to the unmarshaler function.
+func RegisterCustomUnmarshalerContext[T any](unmarshaler func(context.Context, *T, []byte) error) {
+	globalCustomUnmarshalerMu.Lock()
+	defer globalCustomUnmarshalerMu.Unlock()
+
+	var typ *T
+	globalCustomUnmarshalerMap[reflect.TypeOf(typ)] = func(ctx context.Context, v interface{}, b []byte) error {
+		return unmarshaler(ctx, v.(*T), b)
+	}
+}