Add additional support for array indexing in Set

Author: chilland

parser.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"math"
 	"strconv"
-	"strings"
 )
 
 // Errors
@@ -564,23 +563,42 @@ var (
 
 func createInsertComponent(keys []string, setValue []byte, comma, object bool) []byte {
 	var buffer bytes.Buffer
+	isIndex := string(keys[0][0]) == "["
 	if comma {
 		buffer.WriteString(",")
 	}
-	if object {
-		buffer.WriteString("{")
+	if isIndex {
+		buffer.WriteString("[")
+	} else {
+		if object {
+			buffer.WriteString("{")
+		}
+		buffer.WriteString("\"")
+		buffer.WriteString(keys[0])
+		buffer.WriteString("\":")
 	}
-	buffer.WriteString("\"")
-	buffer.WriteString(keys[0])
-	buffer.WriteString("\":")
+
 	for i := 1; i < len(keys); i++ {
-		buffer.WriteString("{\"")
-		buffer.WriteString(keys[i])
-		buffer.WriteString("\":")
+		if string(keys[i][0]) == "[" {
+			buffer.WriteString("[")
+		} else {
+			buffer.WriteString("{\"")
+			buffer.WriteString(keys[i])
+			buffer.WriteString("\":")
+		}
 	}
 	buffer.Write(setValue)
-	buffer.WriteString(strings.Repeat("}", len(keys)-1))
-	if object {
+	for i := len(keys) - 1; i > 0; i-- {
+		if string(keys[i][0]) == "[" {
+			buffer.WriteString("]")
+		} else {
+			buffer.WriteString("}")
+		}
+	}
+	if isIndex {
+		buffer.WriteString("]")
+	}
+	if object && !isIndex {
 		buffer.WriteString("}")
 	}
 	return buffer.Bytes()

parser_test.go

@@ -328,6 +328,46 @@ var setTests = []SetTest{
 		isFound: true,
 		data:    `{"a":"b":"d"}`,
 	},
+	{
+		desc:    "set indexed path to object on empty JSON",
+		json:    `{}`,
+		path:    []string{"top", "[0]", "middle", "[0]", "bottom"},
+		setData: `"value"`,
+		isFound: true,
+		data:    `{"top":[{"middle":[{"bottom":"value"}]}]}`,
+	},
+	{
+		desc:    "set indexed path on existing object with object",
+		json:    `{"top":[{"middle":[]}]}`,
+		path:    []string{"top", "[0]", "middle", "[0]", "bottom"},
+		setData: `"value"`,
+		isFound: true,
+		data:    `{"top":[{"middle":[{"bottom":"value"}]}]}`,
+	},
+	{
+		desc:    "set indexed path on existing object with value",
+		json:    `{"top":[{"middle":[]}]}`,
+		path:    []string{"top", "[0]", "middle", "[0]"},
+		setData: `"value"`,
+		isFound: true,
+		data:    `{"top":[{"middle":["value"]}]}`,
+	},
+	{
+		desc:    "set indexed path on empty object with value",
+		json:    `{}`,
+		path:    []string{"top", "[0]", "middle", "[0]"},
+		setData: `"value"`,
+		isFound: true,
+		data:    `{"top":[{"middle":["value"]}]}`,
+	},
+	{
+		desc:    "set indexed path on object with existing array",
+		json:    `{"top":["one", "two", "three"]}`,
+		path:    []string{"top", "[2]"},
+		setData: `"value"`,
+		isFound: true,
+		data:    `{"top":["one", "two", "value"]}`,
+	},
 }
 
 var getTests = []GetTest{