Jelajahi Sumber

write map with int key

Tao Wen 8 tahun lalu
induk
melakukan
29dc1d407d
2 mengubah file dengan 50 tambahan dan 2 penghapusan
  1. 42 2
      feature_reflect_map.go
  2. 8 0
      jsoniter_map_test.go

+ 42 - 2
feature_reflect_map.go

@@ -3,6 +3,8 @@ package jsoniter
 import (
 	"unsafe"
 	"reflect"
+	"encoding/json"
+	"encoding"
 )
 
 type mapDecoder struct {
@@ -47,13 +49,45 @@ func (encoder *mapEncoder) encode(ptr unsafe.Pointer, stream *Stream) {
 		if i != 0 {
 			stream.WriteMore()
 		}
-		stream.WriteObjectField(key.String())
+		encodeMapKey(key, stream)
+		stream.writeByte(':')
 		val := realVal.MapIndex(key).Interface()
 		encoder.elemEncoder.encodeInterface(val, stream)
 	}
 	stream.WriteObjectEnd()
 }
 
+func encodeMapKey(key reflect.Value, stream *Stream) {
+	if key.Kind() == reflect.String {
+		stream.WriteString(key.String())
+		return
+	}
+	if tm, ok := key.Interface().(encoding.TextMarshaler); ok {
+		buf, err := tm.MarshalText()
+		if err != nil {
+			stream.Error = err
+			return
+		}
+		stream.writeByte('"')
+		stream.Write(buf)
+		stream.writeByte('"')
+		return
+	}
+	switch key.Kind() {
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		stream.writeByte('"')
+		stream.WriteInt64(key.Int())
+		stream.writeByte('"')
+		return
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+		stream.writeByte('"')
+		stream.WriteUint64(key.Uint())
+		stream.writeByte('"')
+		return
+	}
+	stream.Error = &json.UnsupportedTypeError{key.Type()}
+}
+
 func (encoder *mapEncoder) encodeInterface(val interface{}, stream *Stream) {
 	writeToStream(val, stream, encoder)
 }
@@ -84,7 +118,13 @@ func (encoder *mapInterfaceEncoder) encode(ptr unsafe.Pointer, stream *Stream) {
 		if i != 0 {
 			stream.WriteMore()
 		}
-		stream.WriteObjectField(key.String())
+		switch key.Interface().(type) {
+		case string:
+			stream.WriteObjectField(key.String())
+		default:
+			stream.Error = &json.UnsupportedTypeError{key.Type()}
+			return
+		}
 		val := realVal.MapIndex(key).Interface()
 		encoder.elemEncoder.encode(unsafe.Pointer(&val), stream)
 	}

+ 8 - 0
jsoniter_map_test.go

@@ -66,4 +66,12 @@ func Test_slice_of_map(t *testing.T) {
 	val = []map[string]string{}
 	should.Nil(UnmarshalFromString(str, &val))
 	should.Equal("2", val[0]["1"])
+}
+
+func Test_write_int_key_map(t *testing.T) {
+	should := require.New(t)
+	val := map[int]string{1: "2"}
+	str, err := MarshalToString(val)
+	should.Nil(err)
+	should.Equal(`{"1":"2"}`, str)
 }