From 45bdea2a2718d7146892feebcbefb749a5c57421 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Mon, 19 Dec 2022 13:48:03 +0900
Subject: [PATCH] fix custom marshaler (#333)

---
 encode.go      |  2 +-
 encode_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 go.sum         |  1 -
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/encode.go b/encode.go
index 9536eeb..4466abd 100644
--- a/encode.go
+++ b/encode.go
@@ -642,7 +642,7 @@ func (e *Encoder) encodeStruct(ctx context.Context, value reflect.Value, column
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to encode value")
 		}
-		if _, ok := value.(*ast.MappingNode); ok {
+		if _, ok := value.(ast.MapNode); ok {
 			value.AddColumn(e.indent)
 		}
 		var key ast.MapKeyNode = e.encodeString(structField.RenderName, column)
diff --git a/encode_test.go b/encode_test.go
index cfc3d8e..db36028 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -1473,3 +1473,45 @@ a:
 		t.Fatalf("failed to encode. expected %s but got %s", expected, got)
 	}
 }
+
+type customMapSliceOneItemMarshaler struct{}
+
+func (m *customMapSliceOneItemMarshaler) MarshalYAML() ([]byte, error) {
+	var v yaml.MapSlice
+	v = append(v, yaml.MapItem{"a", "b"})
+	return yaml.Marshal(v)
+}
+
+type customMapSliceTwoItemMarshaler struct{}
+
+func (m *customMapSliceTwoItemMarshaler) MarshalYAML() ([]byte, error) {
+	var v yaml.MapSlice
+	v = append(v, yaml.MapItem{"a", "b"})
+	v = append(v, yaml.MapItem{"b", "c"})
+	return yaml.Marshal(v)
+}
+
+func TestCustomMapSliceMarshaler(t *testing.T) {
+	type T struct {
+		A *customMapSliceOneItemMarshaler `yaml:"a"`
+		B *customMapSliceTwoItemMarshaler `yaml:"b"`
+	}
+	b, err := yaml.Marshal(&T{
+		A: &customMapSliceOneItemMarshaler{},
+		B: &customMapSliceTwoItemMarshaler{},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := `
+a:
+  a: b
+b:
+  a: b
+  b: c
+`
+	got := "\n" + string(b)
+	if expected != got {
+		t.Fatalf("failed to encode. expected %s but got %s", expected, got)
+	}
+}
diff --git a/go.sum b/go.sum
index 112c438..fbb7f81 100644
--- a/go.sum
+++ b/go.sum
@@ -28,7 +28,6 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae h1:/WDfKMnPU+m5M4xB+6x4kaepxRw6jWvR5iDRdvjHgy8=
 golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20220406163625-3f8b81556e12 h1:QyVthZKMsyaQwBTJE04jdNN0Pp5Fn9Qga0mrgxyERQM=
 golang.org/x/sys v0.0.0-20220406163625-3f8b81556e12/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=