Merge pull request #231 from kitagry/support-defined-type-key

Support to encode map which has defined type key

encode.go

@@ -432,11 +432,13 @@ func (e *Encoder) encodeMapSlice(ctx context.Context, value MapSlice, column int
 
 func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int) ast.Node {
 	node := ast.Mapping(token.New("", "", e.pos(column)), e.isFlowStyle)
-	keys := []string{}
-	for _, k := range value.MapKeys() {
-		keys = append(keys, k.Interface().(string))
+	keys := make([]interface{}, len(value.MapKeys()))
+	for i, k := range value.MapKeys() {
+		keys[i] = k.Interface()
 	}
-	sort.Strings(keys)
+	sort.Slice(keys, func(i, j int) bool {
+		return fmt.Sprint(keys[i]) < fmt.Sprint(keys[j])
+	})
 	for _, key := range keys {
 		k := reflect.ValueOf(key)
 		v := value.MapIndex(k)
@@ -449,7 +451,7 @@ func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int
 		}
 		node.Values = append(node.Values, ast.MappingValue(
 			nil,
-			e.encodeString(k.Interface().(string), column),
+			e.encodeString(fmt.Sprint(key), column),
 			value,
 		))
 	}

encode_test.go

@@ -601,6 +601,24 @@ func TestEncodeStructIncludeMap(t *testing.T) {
 	}
 }
 
+func TestEncodeDefinedTypeKeyMap(t *testing.T) {
+	type K string
+	type U struct {
+		M map[K]string
+	}
+	bytes, err := yaml.Marshal(U{
+		M: map[K]string{K("x"): "y"},
+	})
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	expect := "m:\n  x: y\n"
+	actual := string(bytes)
+	if actual != expect {
+		t.Fatalf("unexpected output. expect:[%s] actual:[%s]", expect, actual)
+	}
+}
+
 func TestEncodeWithAnchorAndAlias(t *testing.T) {
 	var buf bytes.Buffer
 	enc := yaml.NewEncoder(&buf)