Skip to content

Instantly share code, notes, and snippets.

@Ja7ad
Created August 10, 2024 07:06
Show Gist options
  • Save Ja7ad/8fd3b33518f571fbe82caa50ec90385f to your computer and use it in GitHub Desktop.
Save Ja7ad/8fd3b33518f571fbe82caa50ec90385f to your computer and use it in GitHub Desktop.
PickByNestedKey
package utils
import (
"strconv"
"strings"
)
// PickByNestedKey find a value in nested map
// for nested key need set every key in one string and seperated with dot.
// For example "foo.bar.x.y.z" this is mean you need z key value.
func PickByNestedKey[K comparable, V any](in map[K]V, key string) any {
if key == "" {
return nil
}
keys := strings.Split(key, ".")
var traverse func(data any, keyParts []string) any
traverse = func(data any, keyParts []string) any {
if len(keyParts) == 0 {
return data
}
currentKey := keyParts[0]
remainingKeys := keyParts[1:]
switch d := data.(type) {
case map[K]V:
var key K
if keyStr, ok := any(currentKey).(K); ok {
key = keyStr
} else {
return nil
}
if val, ok := d[key]; ok {
return traverse(val, remainingKeys)
}
case map[string]any:
if val, ok := d[currentKey]; ok {
return traverse(val, remainingKeys)
}
case map[int]any:
idx, err := strconv.Atoi(currentKey)
if err != nil {
return nil
}
if val, ok := d[idx]; ok {
return traverse(val, remainingKeys)
}
case []any:
idx, err := strconv.Atoi(currentKey)
if err != nil || idx < 0 || idx >= len(d) {
return nil
}
return traverse(d[idx], remainingKeys)
}
return nil
}
return traverse(in, keys)
}
package utils
import (
"github.com/stretchr/testify/assert"
"testing"
)
var nestedMapForTest = map[string]any{
"foo": map[string]any{
"bar": map[string]any{
"x": map[string]any{
"y": map[string]any{
"z": "found me",
"a": []any{
map[string]any{
"title": "Calligraphy",
"domain": "art",
},
map[string]any{
"title": "Paint",
"domain": "art",
},
},
},
},
},
},
}
func Test_PickByNestedKey(t *testing.T) {
tests := []struct {
name string
inputMap map[string]any
key string
expected any
}{
{
name: "Full nested key",
inputMap: nestedMapForTest,
key: "foo.bar.x.y.z",
expected: "found me",
},
{
name: "Get Calligraphy",
inputMap: nestedMapForTest,
key: "foo.bar.x.y.a.0.title",
expected: "Calligraphy",
},
{
name: "Top level key",
inputMap: nestedMapForTest,
key: "foo",
expected: map[string]any{"bar": map[string]any{"x": map[string]any{"y": map[string]any{"z": "found me"}}}},
},
{
name: "Mid-level key",
inputMap: nestedMapForTest,
key: "foo.bar",
expected: map[string]any{"x": map[string]any{"y": map[string]any{"z": "found me"}}},
},
{
name: "Non-existent key path",
inputMap: nestedMapForTest,
key: "foo,bar,x,y,z,a,b,c",
expected: nil,
},
{
name: "Invalid key with leading dot",
inputMap: nestedMapForTest,
key: ".foo",
expected: nil,
},
{
name: "Empty key",
inputMap: nestedMapForTest,
key: "",
expected: nil,
},
{
name: "Key not present",
inputMap: nestedMapForTest,
key: "baz",
expected: nil,
},
{
name: "Deep nested but wrong",
inputMap: nestedMapForTest,
key: "foo.bar.x.y.z.b",
expected: nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := PickByNestedKey(tt.inputMap, tt.key)
assert.Equal(t, tt.expected, result)
})
}
}
func Benchmark_PickByNestedKey(b *testing.B) {
keys := []string{
"foo.bar.x.y.z",
"foo.bar.x.y.a.0.title",
"foo.bar.x.y.a.0.domain",
"foo.bar.x.y.a.1.title",
"foo.bar.x.y.a.1.domain",
"foo.bar.x.y",
"foo.bar",
"foo",
}
for _, key := range keys {
b.Run(key, func(b *testing.B) {
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
v := PickByNestedKey(nestedMapForTest, key)
if v == nil {
b.Fatal(v)
}
}
})
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment