Use reflection to allow for any type of maps

walkmap.go

@@ -1,6 +1,8 @@
 package walkmap
 
-import "reflect"
+import (
+	"reflect"
+)
 
 // Visitor is a function that will be called on each non-map node
 // keyPath is a slice containing consecutive keys used to arrive at the given value
@@ -9,19 +11,27 @@ import "reflect"
 type Visitor func(keyPath []interface{}, value interface{}, kind reflect.Kind)
 
 // Walk walks the given map `m` and calls the visitor at every non-map node
-func Walk(m map[interface{}]interface{}, visitor Visitor) {
+func Walk(m interface{}, visitor Visitor) {
+	if reflect.TypeOf(m).Kind() != reflect.Map {
+		panic("m must be a map")
+	}
+
 	walk([]interface{}{}, m, visitor)
 }
 
-func walk(keyPath []interface{}, m map[interface{}]interface{}, visitor Visitor) {
-	for key, val := range m {
+func walk(keyPath []interface{}, m interface{}, visitor Visitor) {
+	mapVal := reflect.ValueOf(m)
+	for _, refKeyValue := range mapVal.MapKeys() {
+		val := mapVal.MapIndex(refKeyValue).Interface()
+		key := refKeyValue.Interface()
+
 		kind := reflect.TypeOf(val).Kind()
 		path := make([]interface{}, len(keyPath))
 		copy(path, keyPath)
 		path = append(path, key)
 
 		if kind == reflect.Map {
-			walk(path, val.(map[interface{}]interface{}), visitor)
+			walk(path, val, visitor)
 		} else {
 			visitor(path, val, kind)
 		}

walkmap_test.go

@@ -137,3 +137,53 @@ func TestWalkNested(t *testing.T) {
 		}
 	}
 }
+
+func TestWalkNestedMixedTypes(t *testing.T) {
+	data := map[interface{}]interface{}{
+		"foo": "1",
+		"bar": "2",
+		"baz": "3",
+		"nested": map[string]interface{}{
+			"fooNested": "4",
+			"barNested": "5",
+			"bazNested": "6",
+			"evenMoreNested": map[interface{}]interface{}{
+				"fooEvenMoreNested": "7",
+				"barEvenMoreNested": "8",
+				"bazEvenMoreNested": "9",
+			},
+		},
+	}
+
+	got := map[string]interface{}{}
+
+	concatenatePath := func(paths []interface{}) string {
+		strPaths := make([]string, 0, len(paths))
+
+		for _, path := range paths {
+			strPaths = append(strPaths, path.(string))
+		}
+
+		return strings.Join(strPaths, ".")
+	}
+
+	Walk(data, func(keyPath []interface{}, value interface{}, kind reflect.Kind) {
+		pathStr := concatenatePath(keyPath)
+		got[pathStr] = value
+	})
+
+	if len(got) != 9 {
+		t.Fatalf("Expected outmap to be of size %d, got %d", 9, len(got))
+	}
+
+	for key, val := range got {
+		exp, err := jsonpath.Get(key, data)
+		if err != nil {
+			t.Fatalf("Expected err to be nil, got: %s", err)
+		}
+
+		if exp != val {
+			t.Fatalf("Expected %s, got %s", exp, val)
+		}
+	}
+}