support unflatten arrays in flatten result

flat.go

@@ -105,7 +105,35 @@ func Unflatten(flat map[string]interface{}, opts *Options) (nested map[string]in
 			Delimiter: ".",
 		}
 	}
-	nested, err = unflatten(flat, opts)
+	if opts.Safe {
+		nested, err = unflatten(flat, opts)
+		return
+	}
+
+	root := &TrieNode{}
+
+	for k, v := range flat {
+		if opts.Prefix != "" {
+			k = strings.TrimPrefix(k, opts.Prefix+opts.Delimiter)
+		}
+
+		// flatten again if value is map[string]interface
+		switch nested := v.(type) {
+		case map[string]interface{}:
+			nested, err := Flatten(v.(map[string]interface{}), opts)
+			if err != nil {
+				return nil, err
+			}
+			parts := strings.Split(k, opts.Delimiter)
+			for newK, newV := range nested {
+				root.insert(append(parts, strings.Split(newK, opts.Delimiter)...), newV)
+			}
+		default:
+			root.insert(strings.Split(k, opts.Delimiter), v)
+		}
+	}
+
+	nested = root.unflatten()
 	return
 }
 

flat_test.go

@@ -379,13 +379,29 @@ func TestUnflatten(t *testing.T) {
 		// },
 	}
 	for i, test := range tests {
-		got, err := Unflatten(test.flat, test.options)
+		opts := test.options
+		got, err := Unflatten(test.flat, opts)
 		if err != nil {
 			t.Errorf("%d: failed to unflatten: %v", i+1, err)
 		}
 		if !reflect.DeepEqual(got, test.want) {
 			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.want)
 		}
+
+		// test safe option
+		if opts == nil {
+			opts = &Options{Delimiter: "."}
+		}
+		opts.Safe = true
+
+		got, err = Unflatten(test.flat, opts)
+		if err != nil {
+			t.Errorf("%d: failed to unflatten with safe option: %v", i+1, err)
+		}
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("%d: mismatch with safe option, got: %v want: %v", i+1, got, test.want)
+		}
+
 	}
 }
 
@@ -437,6 +453,7 @@ func TestFlattenPrefix(t *testing.T) {
 		if err != nil {
 			t.Errorf("%d: failed to unmarshal test: %v", i+1, err)
 		}
+
 		got, err := Flatten(given.(map[string]interface{}), test.options)
 		if err != nil {
 			t.Errorf("%d: failed to flatten: %v", i+1, err)
@@ -499,12 +516,92 @@ func TestUnflattenPrefix(t *testing.T) {
 		},
 	}
 	for i, test := range tests {
-		got, err := Unflatten(test.flat, test.options)
+		opts := test.options
+		got, err := Unflatten(test.flat, opts)
 		if err != nil {
 			t.Errorf("%d: failed to unflatten: %v", i+1, err)
 		}
 		if !reflect.DeepEqual(got, test.want) {
 			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.want)
 		}
+
+		opts.Safe = true
+		got, err = Unflatten(test.flat, opts)
+		if err != nil {
+			t.Errorf("%d: failed to unflatten with safe option: %v", i+1, err)
+		}
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("%d: mismatch with safe option, got: %v want: %v", i+1, got, test.want)
+		}
+	}
+}
+
+type compSlice struct {
+	str  string
+	list []interface{}
+}
+
+func TestSlice(t *testing.T) {
+	tests := []struct {
+		options *Options
+		data    map[string]interface{}
+	}{
+		{nil, map[string]interface{}{"slice": []interface{}{}}},
+		{nil, map[string]interface{}{"slice": []interface{}{1, 2}}},
+		{nil, map[string]interface{}{"slice": []interface{}{"1", "2"}}},
+		{nil, map[string]interface{}{"k": "v", "slice": []interface{}{"1", "2"}}},
+		{nil, map[string]interface{}{"k": "v", "slice": []map[string]string{
+			{"k1": "v1"},
+			{"k2": "v2"},
+		}}},
+		{nil, map[string]interface{}{"k": "v", "slice": [][]string{
+			[]string{"k1", "v1"},
+			[]string{"k2", "v2"},
+		}}},
+		{nil, map[string]interface{}{"k": "v", "slice": []map[string][]string{
+			map[string][]string{"k1": []string{"v11", "v12"}},
+			map[string][]string{"k2": []string{"v21", "v22"}},
+		}}},
+		{nil, map[string]interface{}{"k": "v", "slice": []compSlice{
+			{"k1", []interface{}{1, 2}},
+			{"k2", []interface{}{3, 4}},
+		}}},
+		{nil, map[string]interface{}{"k": "v", "slice": [][]compSlice{
+			{{"k11", []interface{}{1, 2}}, {"k12", []interface{}{3, 4}}},
+			{{"k21", []interface{}{11, 12}}, {"k22", []interface{}{13, 14}}},
+		}}},
+		{nil, map[string]interface{}{"k": "v", "slice": []compSlice{
+			{"k1", []interface{}{
+				compSlice{"k11", []interface{}{1, 2}},
+				compSlice{"k12", []interface{}{3, 4}},
+			}},
+			{"k2", []interface{}{
+				compSlice{"k21", []interface{}{11, 12}},
+				compSlice{"k22", []interface{}{13, 14}},
+			}},
+		}}},
+		{&Options{Prefix: "json", Delimiter: "."}, map[string]interface{}{"k": "v", "slice": []compSlice{
+			{"k1", []interface{}{
+				compSlice{"k11", []interface{}{1, 2}},
+				compSlice{"k12", []interface{}{3, 4}},
+			}},
+			{"k2", []interface{}{
+				compSlice{"k21", []interface{}{11, 12}},
+				compSlice{"k22", []interface{}{13, 14}},
+			}},
+		}}},
+	}
+	for i, test := range tests {
+		f, err := Flatten(test.data, test.options)
+		if err != nil {
+			t.Errorf("%d: failed to flatten: %v", i+1, err)
+		}
+		got, err := Unflatten(f, test.options)
+		if err != nil {
+			t.Errorf("%d: failed to unflatten: %v", i+1, err)
+		}
+		if !reflect.DeepEqual(got, test.data) {
+			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.data)
+		}
 	}
 }

go.mod

@@ -1,4 +1,4 @@
-module github.com/nqd/flat
+module github.com/leslie-qiwa/flat
 
 go 1.16
 

trie.go

@@ -0,0 +1,93 @@
+package flat
+
+import (
+	"fmt"
+	"strconv"
+)
+
+type TrieNode struct {
+	children   map[string]*TrieNode
+	sliceDepth int // record depth of slice in case it is slice of slice
+	value      interface{}
+}
+
+func (t *TrieNode) Print(delimiter string) {
+	t.print("", delimiter)
+}
+
+func (t *TrieNode) print(parent, delimiter string) {
+	if t.value != nil {
+		fmt.Printf("%s: %v\n", parent, t.value)
+	}
+	for k, child := range t.children {
+		if parent != "" {
+			child.print(parent+delimiter+k, delimiter)
+		} else {
+			child.print(k, delimiter)
+		}
+	}
+}
+
+func (t *TrieNode) insert(parts []string, value interface{}) {
+	node := t
+	for i, part := range parts {
+		if node.children == nil {
+			node.children = make(map[string]*TrieNode)
+		}
+		cnode, ok := node.children[part]
+		if !ok {
+			cnode = &TrieNode{}
+			node.children[part] = cnode
+		}
+		node = cnode
+		if i == len(parts)-1 {
+			node.value = value
+		}
+	}
+}
+
+// Start from the bottom to handle slice case
+// TODO: support slice of slice when flatten support it
+func (t *TrieNode) unflatten() map[string]interface{} {
+	ret := make(map[string]interface{})
+	for k, child := range t.children {
+		ret[k] = child.uf()
+	}
+	return ret
+}
+
+func (t *TrieNode) uf() interface{} {
+	if t.value != nil || len(t.children) == 0 {
+		return t.value
+	}
+	isSlice := true
+	sChildren := make([]*TrieNode, len(t.children))
+	for k, v := range t.children {
+		idx, err := strconv.Atoi(k)
+		if err != nil {
+			break
+		}
+		sChildren[idx] = v
+	}
+	for _, v := range sChildren {
+		if v != nil {
+			continue
+		}
+		isSlice = false
+		break
+	}
+
+	if isSlice {
+		ret := make([]interface{}, len(sChildren))
+		for i, child := range sChildren {
+			ret[i] = child.uf()
+		}
+		return ret
+	}
+
+	ret := make(map[string]interface{})
+	for k, child := range t.children {
+		ret[k] = child.uf()
+	}
+	return ret
+}

trie_test.go

@@ -0,0 +1,108 @@
+package flat
+
+import (
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestTrie(t *testing.T) {
+	tests := []struct {
+		data map[string]interface{}
+		want map[string]interface{}
+	}{
+		{
+			map[string]interface{}{"hello": "world"},
+			map[string]interface{}{"hello": "world"},
+		},
+		{
+			map[string]interface{}{"hello.world.again": "good morning"},
+			map[string]interface{}{
+				"hello": map[string]interface{}{
+					"world": map[string]interface{}{
+						"again": "good morning",
+					},
+				},
+			},
+		},
+		{
+			map[string]interface{}{"a.0.0": 1, "a.0.1": 2, "a.1.0": 21, "a.1.1": 22},
+			map[string]interface{}{
+				"a": []interface{}{
+					[]interface{}{1, 2},
+					[]interface{}{21, 22},
+				},
+			},
+		},
+		{
+			map[string]interface{}{"a.0.0": "1", "a.0.1": "2", "a.1.0": "21", "a.1.1": "22"},
+			map[string]interface{}{
+				"a": []interface{}{
+					[]interface{}{"1", "2"},
+					[]interface{}{"21", "22"},
+				},
+			},
+		},
+		{
+			map[string]interface{}{"a.0": "1", "a.1": "2", "b": "21"},
+			map[string]interface{}{"a": []interface{}{"1", "2"}, "b": "21"},
+		},
+		{
+			map[string]interface{}{"a.b.0": "1", "a.b.1": "2", "c": "21"},
+			map[string]interface{}{"a": map[string]interface{}{"b": []interface{}{"1", "2"}}, "c": "21"},
+		},
+		{
+			map[string]interface{}{"a.0.b.0": "1", "a.0.b.1": "2", "a.1.b.0": "3", "a.1.b.1": "4", "c": "21"},
+			map[string]interface{}{
+				"a": []interface{}{
+					map[string]interface{}{
+						"b": []interface{}{"1", "2"},
+					},
+					map[string]interface{}{
+						"b": []interface{}{"3", "4"},
+					},
+				},
+				"c": "21",
+			},
+		},
+		{
+			map[string]interface{}{
+				"a.0.b.0.d": "1", "a.0.b.0.e": "2",
+				"a.0.b.1.d": "3", "a.0.b.1.e": "4",
+				"a.1.b.0.d": "11", "a.1.b.0.e": "12",
+				"a.1.b.0.f": "13", "a.1.b.0.g": "14",
+				"a.1.b.1.d": "15", "a.1.b.1.e": "16",
+				"a.1.b.1.f": "17", "a.1.b.1.g": "18",
+				"c": "21"},
+			map[string]interface{}{
+				"a": []interface{}{
+					map[string]interface{}{
+						"b": []interface{}{
+							map[string]interface{}{"d": "1", "e": "2"},
+							map[string]interface{}{"d": "3", "e": "4"},
+						},
+					},
+					map[string]interface{}{
+						"b": []interface{}{
+							map[string]interface{}{"d": "11", "e": "12", "f": "13", "g": "14"},
+							map[string]interface{}{"d": "15", "e": "16", "f": "17", "g": "18"},
+						},
+					},
+				},
+				"c": "21",
+			},
+		},
+	}
+
+	for i, test := range tests {
+		root := &TrieNode{}
+		for k, v := range test.data {
+			root.insert(strings.Split(k, "."), v)
+		}
+
+		got := root.unflatten()
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.want)
+		}
+	}
+}