feat: ✨ unmarshal into ordered maps

json.go

@@ -20,7 +20,7 @@ var (
 // MarshalJSON implements the json.Marshaler interface.
 func (om *OrderedMap[K, V]) MarshalJSON() ([]byte, error) { //nolint:funlen
 	if om == nil || om.list == nil {
-		return []byte("null"), nil
+		return JSONNullBytes(), nil
 	}
 
 	writer := jwriter.Writer{
@@ -129,7 +129,6 @@ func (om *OrderedMap[K, V]) UnmarshalJSON(data []byte) error {
 			}
 
 			var key K
-			var value V
 
 			switch typedKey := any(&key).(type) {
 			case *string:
@@ -168,6 +167,29 @@ func (om *OrderedMap[K, V]) UnmarshalJSON(data []byte) error {
 				}
 			}
 
+			tValue := reflect.TypeFor[V]()
+			tValueKind := tValue.Kind()
+
+			var value V
+
+			// in the case of an `any`
+			// the value may be nil, we should represent it as such instead of an empty map
+			// otherwise
+			// try to unmarshall into an ordered map before falling back to
+			// what json.Unmarshal would regularly do, which, when the value is a map, it would use map[string]any
+			if tValueKind == reflect.Interface {
+				if bytes.Equal(valueData, JSONNullBytes()) {
+					om.Set(key, value)
+					return nil
+				}
+				om1 := New[string, any]()
+				if err := json.Unmarshal(valueData, &om1); err == nil {
+					om.Set(key, any(om1).(V))
+					return nil
+				}
+			}
+
+			// fallback to json.Unmarshal of V
 			if err := json.Unmarshal(valueData, &value); err != nil {
 				return err
 			}

json_test.go

@@ -202,6 +202,21 @@ func TestUnmarshallJSON(t *testing.T) {
 
 		assertLenEqual(t, om, 0)
 	})
+
+	t.Run("nested maps become ordered", func(t *testing.T) {
+		data := `{"foo":{"zebra":"z","apple":"a"}}`
+
+		om := New[string, any]()
+		require.NoError(t, json.Unmarshal([]byte(data), &om))
+
+		expected := New[string, any]()
+		expected.Set("zebra", "z")
+		expected.Set("apple", "a")
+
+		got, ok := om.Get("foo")
+		require.True(t, ok)
+		assert.Equal(t, expected, got)
+	})
 }
 
 // const specialCharacters = "\\\\/\"\b\f\n\r\t\x00\uffff\ufffd世界\u007f\u00ff\U0010FFFF"

null.go

@@ -0,0 +1,5 @@
+package orderedmap
+
+func JSONNullBytes() []byte {
+	return []byte("null")
+}

yaml.go

@@ -1,7 +1,10 @@
 package orderedmap
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
+	"reflect"
 
 	"gopkg.in/yaml.v3"
 )
@@ -14,7 +17,7 @@ var (
 // MarshalYAML implements the yaml.Marshaler interface.
 func (om *OrderedMap[K, V]) MarshalYAML() (interface{}, error) {
 	if om == nil {
-		return []byte("null"), nil
+		return JSONNullBytes(), nil
 	}
 
 	node := yaml.Node{
@@ -55,11 +58,29 @@ func (om *OrderedMap[K, V]) UnmarshalYAML(value *yaml.Node) error {
 
 	for index := 0; index < len(value.Content); index += 2 {
 		var key K
-		var val V
-
 		if err := value.Content[index].Decode(&key); err != nil {
 			return err
 		}
+
+		tValue := reflect.TypeFor[V]()
+		tValueKind := tValue.Kind()
+
+		var val V
+
+		if tValueKind == reflect.Interface {
+			valueNode := value.Content[index+1]
+			if bytes.Equal([]byte(valueNode.Value), JSONNullBytes()) {
+				om.Set(key, val)
+				continue
+			}
+
+			om1, err := tryDecodeToOrderedMap(value.Content[index+1].Decode)
+			if err == nil {
+				om.Set(key, om1.(V))
+				continue
+			}
+		}
+
 		if err := value.Content[index+1].Decode(&val); err != nil {
 			return err
 		}
@@ -69,3 +90,27 @@ func (om *OrderedMap[K, V]) UnmarshalYAML(value *yaml.Node) error {
 
 	return nil
 }
+
+// since yaml supports integers and strings as map keys, we need to try both
+func tryDecodeToOrderedMap(decoderFunc func(v any) (err error)) (any, error) {
+	type newFunc func() any
+
+	possibilities := []newFunc{
+		func() any { return New[int, any]() },
+		func() any { return New[int32, any]() },
+		func() any { return New[int64, any]() },
+		func() any { return New[uint, any]() },
+		func() any { return New[uint32, any]() },
+		func() any { return New[uint64, any]() },
+		func() any { return New[string, any]() },
+	}
+
+	for _, possibilityFunc := range possibilities {
+		possibility := possibilityFunc()
+		if err := decoderFunc(possibility); err == nil {
+			return possibility, nil
+		}
+	}
+
+	return nil, errors.New("could not decode to an ordered map")
+}

yaml_test.go

@@ -1,10 +1,11 @@
 package orderedmap
 
 import (
+	"testing"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v3"
-	"testing"
 )
 
 func TestMarshalYAML(t *testing.T) {
@@ -202,6 +203,21 @@ func TestUnmarshallYAML(t *testing.T) {
 
 		assertLenEqual(t, om, 0)
 	})
+
+	t.Run("unmarshals nests maps as ordered maps", func(t *testing.T) {
+		data := `{"foo":{"zebra":"z","apple":"a"}}`
+
+		om := New[string, any]()
+		require.NoError(t, yaml.Unmarshal([]byte(data), &om))
+
+		expected := New[string, any]()
+		expected.Set("zebra", "z")
+		expected.Set("apple", "a")
+
+		actual, ok := om.Get("foo")
+		require.True(t, ok)
+		require.Equal(t, expected, actual)
+	})
 }
 
 func TestYAMLSpecialCharacters(t *testing.T) {
@@ -264,6 +280,7 @@ m:
             foo: bar
     foo:
         - 12:
+            z: null
             b: true
             i: 12
             m:
@@ -277,8 +294,8 @@ m:
                 - 2
                 - 3
         - 3:
-            c: null
             d: 87
+            c: null
           4:
             e: true
           5: