Merge pull request #219 from shmel1k/feature/issue_212_implement_key_text_marshaler

Implement encoding.TextMarshaler interface for map keys

Makefile

@@ -26,7 +26,7 @@ generate: root build
 		.root/src/$(PKG)/tests/named_type.go \
 		.root/src/$(PKG)/tests/custom_map_key_type.go \
 		.root/src/$(PKG)/tests/embedded_type.go \
-		.root/src/$(PKG)/tests/reference_to_pointer.go
+		.root/src/$(PKG)/tests/reference_to_pointer.go \
 
 	.root/bin/easyjson -all .root/src/$(PKG)/tests/data.go
 	.root/bin/easyjson -all .root/src/$(PKG)/tests/nothing.go
@@ -39,6 +39,7 @@ generate: root build
 	.root/bin/easyjson .root/src/$(PKG)/tests/custom_map_key_type.go
 	.root/bin/easyjson .root/src/$(PKG)/tests/embedded_type.go
 	.root/bin/easyjson .root/src/$(PKG)/tests/reference_to_pointer.go
+	.root/bin/easyjson .root/src/$(PKG)/tests/key_marshaler_map.go
 	.root/bin/easyjson -disallow_unknown_fields .root/src/$(PKG)/tests/disallow_unknown.go
 
 test: generate root

gen/decoder.go

@@ -240,7 +240,13 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 		fmt.Fprintln(g.out, ws+"  }")
 
 		fmt.Fprintln(g.out, ws+"  for !in.IsDelim('}') {")
-		if keyDec != "" {
+		// NOTE: extra check for TextUnmarshaler. It overrides default methods.
+		if reflect.PtrTo(key).Implements(reflect.TypeOf((*encoding.TextUnmarshaler)(nil)).Elem()) {
+			fmt.Fprintln(g.out, ws+"    var key "+g.getType(key))
+			fmt.Fprintln(g.out, ws+"if data := in.UnsafeBytes(); in.Ok() {")
+			fmt.Fprintln(g.out, ws+"  in.AddError(key.UnmarshalText(data) )")
+			fmt.Fprintln(g.out, ws+"}")
+		} else if keyDec != "" {
 			fmt.Fprintln(g.out, ws+"    key := "+g.getType(key)+"("+keyDec+")")
 		} else {
 			fmt.Fprintln(g.out, ws+"    var key "+g.getType(key))

gen/encoder.go

@@ -223,7 +223,11 @@ func (g *Generator) genTypeEncoderNoCheck(t reflect.Type, in string, tags fieldT
 		fmt.Fprintln(g.out, ws+"  "+tmpVar+"First := true")
 		fmt.Fprintln(g.out, ws+"  for "+tmpVar+"Name, "+tmpVar+"Value := range "+in+" {")
 		fmt.Fprintln(g.out, ws+"    if "+tmpVar+"First { "+tmpVar+"First = false } else { out.RawByte(',') }")
-		if keyEnc != "" {
+
+		// NOTE: extra check for TextMarshaler. It overrides default methods.
+		if reflect.PtrTo(key).Implements(reflect.TypeOf((*encoding.TextMarshaler)(nil)).Elem()) {
+			fmt.Fprintln(g.out, ws+"    "+fmt.Sprintf("out.RawText(("+tmpVar+"Name).MarshalText()"+")"))
+		} else if keyEnc != "" {
 			fmt.Fprintln(g.out, ws+"    "+fmt.Sprintf(keyEnc, tmpVar+"Name"))
 		} else {
 			if err := g.genTypeEncoder(key, tmpVar+"Name", tags, indent+2, false); err != nil {

tests/basic_test.go

@@ -52,6 +52,7 @@ var testCases = []struct {
 	{&intArrayStructValue, intArrayStructValueString},
 	{&myUInt8SliceValue, myUInt8SliceString},
 	{&myUInt8ArrayValue, myUInt8ArrayString},
+	{&mapWithEncodingMarshaler, mapWithEncodingMarshalerString},
 }
 
 func TestMarshal(t *testing.T) {

tests/key_marshaler_map.go

@@ -0,0 +1,20 @@
+package tests
+
+type KeyWithEncodingMarshaler int
+
+func (f KeyWithEncodingMarshaler) MarshalText() (text []byte, err error) {
+	return []byte("hello"), nil
+}
+
+func (f *KeyWithEncodingMarshaler) UnmarshalText(text []byte) error {
+	if string(text) == "hello" {
+		*f = 5
+	}
+	return nil
+}
+
+//easyjson:json
+type KeyWithEncodingMarshalers map[KeyWithEncodingMarshaler]string
+
+var mapWithEncodingMarshaler KeyWithEncodingMarshalers = KeyWithEncodingMarshalers{5: "hello"}
+var mapWithEncodingMarshalerString = `{"hello":"hello"}`