GODRIVER-1494 add option for encoding/json style map encoding (#345)

Author: iwysiu

bson/bson_test.go

@@ -8,11 +8,17 @@ package bson
 
 import (
 	"bytes"
+	"fmt"
+	"reflect"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/bson/bsonoptions"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
@@ -112,6 +118,71 @@ func TestD(t *testing.T) {
 	})
 }
 
+type stringerString string
+
+func (ss stringerString) String() string {
+	return "bar"
+}
+
+type keyBool bool
+
+func (kb keyBool) MarshalKey() (string, error) {
+	return fmt.Sprintf("%v", kb), nil
+}
+
+func (kb keyBool) UnmarshalKey(key string) error {
+	switch key {
+	case "true":
+		kb = true
+	case "false":
+		kb = false
+	default:
+		return fmt.Errorf("invalid bool value %v", key)
+	}
+	return nil
+}
+
+func TestMapCodec(t *testing.T) {
+	t.Run("EncodeKeysWithStringer", func(t *testing.T) {
+		strstr := stringerString("foo")
+		mapObj := map[stringerString]int{strstr: 1}
+		testCases := []struct {
+			name string
+			opts *bsonoptions.MapCodecOptions
+			key  string
+		}{
+			{"default", bsonoptions.MapCodec(), "foo"},
+			{"true", bsonoptions.MapCodec().SetEncodeKeysWithStringer(true), "bar"},
+			{"false", bsonoptions.MapCodec().SetEncodeKeysWithStringer(false), "foo"},
+		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				mapCodec := bsoncodec.NewMapCodec(tc.opts)
+				mapRegistry := NewRegistryBuilder().RegisterDefaultEncoder(reflect.Map, mapCodec).Build()
+				val, err := MarshalWithRegistry(mapRegistry, mapObj)
+				assert.Nil(t, err, "Marshal error: %v", err)
+				assert.True(t, strings.Contains(string(val), tc.key), "expected result to contain %v, got: %v", tc.key, string(val))
+			})
+		}
+	})
+	t.Run("keys implements keyMarshaler and keyUnmarshaler", func(t *testing.T) {
+		mapObj := map[keyBool]int{keyBool(false): 1}
+
+		doc, err := Marshal(mapObj)
+		assert.Nil(t, err, "Marshal error: %v", err)
+		idx, want := bsoncore.AppendDocumentStart(nil)
+		want = bsoncore.AppendInt32Element(want, "false", 1)
+		want, _ = bsoncore.AppendDocumentEnd(want, idx)
+		assert.Equal(t, want, doc, "expected result %v, got %v", string(want), string(doc))
+
+		var got map[keyBool]int
+		err = Unmarshal(doc, &got)
+		assert.Nil(t, err, "Unmarshal error: %v", err)
+		assert.Equal(t, mapObj, got, "expected result %v, got %v", mapObj, got)
+
+	})
+}
+
 func TestExtJSONEscapeKey(t *testing.T) {
 	doc := D{{Key: "\\usb#", Value: int32(1)}}
 	b, err := MarshalExtJSON(&doc, false, false)

bson/bsoncodec/default_value_decoders_test.go

@@ -808,7 +808,7 @@ func TestDefaultValueDecoders(t *testing.T) {
 					&DecodeContext{Registry: buildDefaultRegistry()},
 					&bsonrwtest.ValueReaderWriter{},
 					bsonrwtest.ReadElement,
-					fmt.Errorf("BSON map must have string or decimal keys. Got:%v", reflect.ValueOf(map[bool]interface{}{}).Type()),
+					fmt.Errorf("unsupported key type: %T", false),
 				},
 				{
 					"ReadDocument Error",

bson/bsoncodec/map_codec.go

@@ -20,12 +20,29 @@ var defaultMapCodec = NewMapCodec()
 
 // MapCodec is the Codec used for map values.
 type MapCodec struct {
-	DecodeZerosMap   bool
-	EncodeNilAsEmpty bool
+	DecodeZerosMap         bool
+	EncodeNilAsEmpty       bool
+	EncodeKeysWithStringer bool
 }
 
 var _ ValueCodec = &MapCodec{}
 
+// KeyMarshaler is the interface implemented by an object that can marshal itself into a string key.
+// This applies to types used as map keys and is similar to encoding.TextMarshaler.
+type KeyMarshaler interface {
+	MarshalKey() (key string, err error)
+}
+
+// KeyUnmarshaler is the interface implemented by an object that can unmarshal a string representation
+// of itself. This applies to types used as map keys and is similar to encoding.TextUnmarshaler.
+//
+// UnmarshalKey must be able to decode the form generated by MarshalKey.
+// UnmarshalKey must copy the text if it wishes to retain the text
+// after returning.
+type KeyUnmarshaler interface {
+	UnmarshalKey(key string) error
+}
+
 // NewMapCodec returns a MapCodec with options opts.
 func NewMapCodec(opts ...*bsonoptions.MapCodecOptions) *MapCodec {
 	mapOpt := bsonoptions.MergeMapCodecOptions(opts...)
@@ -37,6 +54,9 @@ func NewMapCodec(opts ...*bsonoptions.MapCodecOptions) *MapCodec {
 	if mapOpt.EncodeNilAsEmpty != nil {
 		codec.EncodeNilAsEmpty = *mapOpt.EncodeNilAsEmpty
 	}
+	if mapOpt.EncodeKeysWithStringer != nil {
+		codec.EncodeKeysWithStringer = *mapOpt.EncodeKeysWithStringer
+	}
 	return &codec
 }
 
@@ -79,7 +99,11 @@ func (mc *MapCodec) mapEncodeValue(ec EncodeContext, dw bsonrw.DocumentWriter, v
 
 	keys := val.MapKeys()
 	for _, key := range keys {
-		keyStr := fmt.Sprint(key)
+		keyStr, err := mc.encodeKey(key)
+		if err != nil {
+			return err
+		}
+
 		if collisionFn != nil && collisionFn(keyStr) {
 			return fmt.Errorf("Key %s of inlined map conflicts with a struct field name", key)
 		}
@@ -160,7 +184,6 @@ func (mc *MapCodec) DecodeValue(dc DecodeContext, vr bsonrw.ValueReader, val ref
 	}
 
 	keyType := val.Type().Key()
-	keyKind := keyType.Kind()
 
 	for {
 		key, vr, err := dr.ReadElement()
@@ -171,23 +194,9 @@ func (mc *MapCodec) DecodeValue(dc DecodeContext, vr bsonrw.ValueReader, val ref
 			return err
 		}
 
-		k := reflect.ValueOf(key)
-		if keyType != tString {
-			switch keyKind {
-			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
-				reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
-				reflect.Float32, reflect.Float64:
-				parsed, err := strconv.ParseFloat(k.String(), 64)
-				if err != nil {
-					return fmt.Errorf("Map key is defined to be a decimal type (%v) but got error %v", keyKind, err)
-				}
-				k = reflect.ValueOf(parsed)
-			case reflect.String: // if keyType wraps string
-			default:
-				return fmt.Errorf("BSON map must have string or decimal keys. Got:%v", val.Type())
-			}
-
-			k = k.Convert(keyType)
+		k, err := mc.decodeKey(key, keyType)
+		if err != nil {
+			return err
 		}
 
 		elem := reflect.New(eType).Elem()
@@ -207,3 +216,82 @@ func clearMap(m reflect.Value) {
 		m.SetMapIndex(k, none)
 	}
 }
+
+func (mc *MapCodec) encodeKey(val reflect.Value) (string, error) {
+	if mc.EncodeKeysWithStringer {
+		return fmt.Sprint(val), nil
+	}
+
+	// keys of any string type are used directly
+	if val.Kind() == reflect.String {
+		return val.String(), nil
+	}
+	// KeyMarshalers are marshaled
+	if km, ok := val.Interface().(KeyMarshaler); ok {
+		if val.Kind() == reflect.Ptr && val.IsNil() {
+			return "", nil
+		}
+		buf, err := km.MarshalKey()
+		if err == nil {
+			return buf, nil
+		}
+		return "", err
+	}
+
+	switch val.Kind() {
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		return strconv.FormatInt(val.Int(), 10), nil
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+		return strconv.FormatUint(val.Uint(), 10), nil
+	}
+	return "", fmt.Errorf("unsupported key type: %v", val.Type())
+}
+
+var keyUnmarshalerType = reflect.TypeOf((*KeyUnmarshaler)(nil)).Elem()
+
+func (mc *MapCodec) decodeKey(key string, keyType reflect.Type) (reflect.Value, error) {
+	keyVal := reflect.ValueOf(key)
+	var err error
+	switch {
+	// First, if EncodeKeysWithStringer is not enabled, try to decode withKeyUnmarshaler
+	case !mc.EncodeKeysWithStringer && reflect.PtrTo(keyType).Implements(keyUnmarshalerType):
+		keyVal = reflect.New(keyType)
+		v := keyVal.Interface().(KeyUnmarshaler)
+		err = v.UnmarshalKey(key)
+		keyVal = keyVal.Elem()
+	// Otherwise, go to type specific behavior
+	default:
+		switch keyType.Kind() {
+		case reflect.String:
+			keyVal = reflect.ValueOf(key).Convert(keyType)
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			s := string(key)
+			n, parseErr := strconv.ParseInt(s, 10, 64)
+			if parseErr != nil || reflect.Zero(keyType).OverflowInt(n) {
+				err = fmt.Errorf("failed to unmarshal number key %v", s)
+			}
+			keyVal = reflect.ValueOf(n).Convert(keyType)
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+			s := string(key)
+			n, parseErr := strconv.ParseUint(s, 10, 64)
+			if parseErr != nil || reflect.Zero(keyType).OverflowUint(n) {
+				err = fmt.Errorf("failed to unmarshal number key %v", s)
+				break
+			}
+			keyVal = reflect.ValueOf(n).Convert(keyType)
+		case reflect.Float32, reflect.Float64:
+			if mc.EncodeKeysWithStringer {
+				parsed, err := strconv.ParseFloat(key, 64)
+				if err != nil {
+					return keyVal, fmt.Errorf("Map key is defined to be a decimal type (%v) but got error %v", keyType.Kind(), err)
+				}
+				keyVal = reflect.ValueOf(parsed)
+				break
+			}
+			fallthrough
+		default:
+			return keyVal, fmt.Errorf("unsupported key type: %v", keyType)
+		}
+	}
+	return keyVal, err
+}

bson/bsonoptions/map_codec_options.go

@@ -10,6 +10,12 @@ package bsonoptions
 type MapCodecOptions struct {
 	DecodeZerosMap   *bool // Specifies if the map should be zeroed before decoding into it. Defaults to false.
 	EncodeNilAsEmpty *bool // Specifies if a nil map should encode as an empty document instead of null. Defaults to false.
+	// Specifies how keys should be handled. If false, the behavior matches encoding/json, where the encoding key type must
+	// either be a string, an integer type, or implement bsoncodec.KeyMarshaler and the decoding key type must either be a
+	// string, an integer type, or implement bsoncodec.KeyUnmarshaler. If true, keys are encoded with fmt.Sprint() and the
+	// encoding key type must be a string, an integer type, or a float. If true, the use of Stringer will override
+	// TextMarshaler/TextUnmarshaler. Defaults to false.
+	EncodeKeysWithStringer *bool
 }
 
 // MapCodec creates a new *MapCodecOptions
@@ -23,12 +29,22 @@ func (t *MapCodecOptions) SetDecodeZerosMap(b bool) *MapCodecOptions {
 	return t
 }
 
-// SetEncodeNilAsEmpty specifies  if a nil map should encode as an empty document instead of null. Defaults to false.
+// SetEncodeNilAsEmpty specifies if a nil map should encode as an empty document instead of null. Defaults to false.
 func (t *MapCodecOptions) SetEncodeNilAsEmpty(b bool) *MapCodecOptions {
 	t.EncodeNilAsEmpty = &b
 	return t
 }
 
+// SetEncodeKeysWithStringer specifies how keys should be handled. If false, the behavior matches encoding/json, where the
+// encoding key type must either be a string, an integer type, or implement bsoncodec.KeyMarshaler and the decoding key
+// type must either be a string, an integer type, or implement bsoncodec.KeyUnmarshaler. If true, keys are encoded with
+// fmt.Sprint() and the encoding key type must be a string, an integer type, or a float. If true, the use of Stringer
+// will override TextMarshaler/TextUnmarshaler. Defaults to false.
+func (t *MapCodecOptions) SetEncodeKeysWithStringer(b bool) *MapCodecOptions {
+	t.EncodeKeysWithStringer = &b
+	return t
+}
+
 // MergeMapCodecOptions combines the given *MapCodecOptions into a single *MapCodecOptions in a last one wins fashion.
 func MergeMapCodecOptions(opts ...*MapCodecOptions) *MapCodecOptions {
 	s := MapCodec()
@@ -42,6 +58,9 @@ func MergeMapCodecOptions(opts ...*MapCodecOptions) *MapCodecOptions {
 		if opt.EncodeNilAsEmpty != nil {
 			s.EncodeNilAsEmpty = opt.EncodeNilAsEmpty
 		}
+		if opt.EncodeKeysWithStringer != nil {
+			s.EncodeKeysWithStringer = opt.EncodeKeysWithStringer
+		}
 	}
 
 	return s

bson/mgocompat/registry.go

@@ -58,7 +58,8 @@ func NewRegistryBuilder() *bsoncodec.RegistryBuilder {
 	mapCodec := bsoncodec.NewMapCodec(
 		bsonoptions.MapCodec().
 			SetDecodeZerosMap(true).
-			SetEncodeNilAsEmpty(true))
+			SetEncodeNilAsEmpty(true).
+			SetEncodeKeysWithStringer(true))
 	uintcodec := bsoncodec.NewUIntCodec(bsonoptions.UIntCodec().SetEncodeToMinSize(true))
 
 	rb.RegisterTypeDecoder(tEmpty, emptyInterCodec).