Refactor write to allow new entries

Author: Mike Farah

data_navigator.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	// "fmt"
 	"github.com/mikefarah/yaml/Godeps/_workspace/src/gopkg.in/yaml.v2"
 	"strconv"
 )
@@ -16,40 +15,90 @@ func entryInSlice(context yaml.MapSlice, key interface{}) *yaml.MapItem {
 	return nil
 }
 
-func write(context yaml.MapSlice, head string, tail []string, value interface{}) {
-	if len(tail) == 0 {
-		var entry = entryInSlice(context, head)
-		entry.Value = value
-	} else {
-		// e.g. if updating a.b.c, we need to get the 'b', this could be a map or an array
-		var parent = readMap(context, head, tail[0:len(tail)-1])
-		switch parent.(type) {
-		case yaml.MapSlice:
-			toUpdate := parent.(yaml.MapSlice)
-			// b is a map, update the key 'c' to the supplied value
-			key := (tail[len(tail)-1])
-			toUpdateEntry := entryInSlice(toUpdate, key)
-			toUpdateEntry.Value = value
-		case []interface{}:
-			toUpdate := parent.([]interface{})
-			// b is an array, update it at index 'c' to the supplied value
-			rawIndex := (tail[len(tail)-1])
-			index, err := strconv.ParseInt(rawIndex, 10, 64)
-			if err != nil {
-				die("Error accessing array: %v", err)
-			}
-			toUpdate[index] = value
-		}
+func writeMap(context interface{}, paths []string, value interface{}) yaml.MapSlice {
+	log.Debugf("writeMap for %v for %v with value %v\n", paths, context, value)
+
+	var mapSlice yaml.MapSlice
+	switch context.(type) {
+	case yaml.MapSlice:
+		mapSlice = context.(yaml.MapSlice)
+	default:
+		mapSlice = make(yaml.MapSlice, 0)
+	}
+
+	if len(paths) == 0 {
+		return mapSlice
+	}
+
+	child := entryInSlice(mapSlice, paths[0])
+	if child == nil {
+		newChild := yaml.MapItem{Key: paths[0]}
+		mapSlice = append(mapSlice, newChild)
+		child = entryInSlice(mapSlice, paths[0])
+		log.Debugf("\tAppended child at %v for mapSlice %v\n", paths[0], mapSlice)
+	}
+
+	log.Debugf("\tchild.Value %v\n", child.Value)
+
+	remainingPaths := paths[1:len(paths)]
+	child.Value = updatedChildValue(child.Value, remainingPaths, value)
+	log.Debugf("\tReturning mapSlice %v\n", mapSlice)
+	return mapSlice
+}
+
+func updatedChildValue(child interface{}, remainingPaths []string, value interface{}) interface{} {
+	if len(remainingPaths) == 0 {
+		return value
+	}
+
+	_, nextIndexErr := strconv.ParseInt(remainingPaths[0], 10, 64)
+	if nextIndexErr != nil {
+		// must be a map
+		return writeMap(child, remainingPaths, value)
+	}
+
+	// must be an array
+	return writeArray(child, remainingPaths, value)
+}
+
+func writeArray(context interface{}, paths []string, value interface{}) []interface{} {
+	log.Debugf("writeArray for %v for %v with value %v\n", paths, context, value)
+	var array []interface{}
+	switch context.(type) {
+	case []interface{}:
+		array = context.([]interface{})
+	default:
+		array = make([]interface{}, 1)
+	}
 
+	if len(paths) == 0 {
+		return array
 	}
+
+	log.Debugf("\tarray %v\n", array)
+
+	rawIndex := paths[0]
+	index, err := strconv.ParseInt(rawIndex, 10, 64)
+	if err != nil {
+		die("Error accessing array: %v", err)
+	}
+	currentChild := array[index]
+
+	log.Debugf("\tcurrentChild %v\n", currentChild)
+
+	remainingPaths := paths[1:len(paths)]
+	array[index] = updatedChildValue(currentChild, remainingPaths, value)
+	log.Debugf("\tReturning array %v\n", array)
+	return array
 }
 
 func readMap(context yaml.MapSlice, head string, tail []string) interface{} {
 	if head == "*" {
 		return readMapSplat(context, tail)
 	}
-	entry := entryInSlice(context, head)
 	var value interface{}
+
+	entry := entryInSlice(context, head)
 	if entry != nil {
 		value = entry.Value
 	}

data_navigator_test.go

@@ -3,10 +3,18 @@ package main
 import (
 	"fmt"
 	"github.com/mikefarah/yaml/Godeps/_workspace/src/gopkg.in/yaml.v2"
+	"github.com/op/go-logging"
+	"os"
 	"sort"
 	"testing"
 )
 
+func TestMain(m *testing.M) {
+	backend.SetLevel(logging.ERROR, "")
+	logging.SetBackend(backend)
+	os.Exit(m.Run())
+}
+
 func TestReadMap_simple(t *testing.T) {
 	var data = parseData(`
 ---
@@ -113,8 +121,8 @@ func TestWrite_really_simple(t *testing.T) {
     b: 2
 `)
 
-	write(data, "b", []string{}, "4")
-	b := entryInSlice(data, "b").Value
+	updated := writeMap(data, []string{"b"}, "4")
+	b := entryInSlice(updated, "b").Value
 	assertResult(t, "4", b)
 }
 
@@ -124,31 +132,86 @@ b:
   c: 2
 `)
 
-	write(data, "b", []string{"c"}, "4")
-	b := entryInSlice(data, "b").Value.(yaml.MapSlice)
+	updated := writeMap(data, []string{"b", "c"}, "4")
+	b := entryInSlice(updated, "b").Value.(yaml.MapSlice)
 	c := entryInSlice(b, "c").Value
 	assertResult(t, "4", c)
 }
 
+func TestWrite_new(t *testing.T) {
+	var data = parseData(`
+b:
+  c: 2
+`)
+
+	updated := writeMap(data, []string{"b", "d"}, "4")
+	b := entryInSlice(updated, "b").Value.(yaml.MapSlice)
+	d := entryInSlice(b, "d").Value
+	assertResult(t, "4", d)
+}
+
+func TestWrite_new_deep(t *testing.T) {
+	var data = parseData(`
+b:
+  c: 2
+`)
+
+	updated := writeMap(data, []string{"b", "d", "f"}, "4")
+	assertResult(t, "4", readMap(updated, "b", []string{"d", "f"}))
+}
+
 func TestWrite_array(t *testing.T) {
 	var data = parseData(`
 b:
   - aa
 `)
 
-	write(data, "b", []string{"0"}, "bb")
+	updated := writeMap(data, []string{"b", "0"}, "bb")
 
-	b := entryInSlice(data, "b").Value.([]interface{})
+	b := entryInSlice(updated, "b").Value.([]interface{})
 	assertResult(t, "bb", b[0].(string))
 }
 
+func TestWrite_new_array(t *testing.T) {
+	var data = parseData(`
+b:
+  c: 2
+`)
+
+	updated := writeMap(data, []string{"b", "0"}, "4")
+	assertResult(t, "4", readMap(updated, "b", []string{"0"}))
+}
+
+func TestWrite_new_array_deep(t *testing.T) {
+	var data = parseData(`
+b:
+  c: 2
+`)
+
+	var expected = `b:
+- c: "4"`
+
+	updated := writeMap(data, []string{"b", "0", "c"}, "4")
+	assertResult(t, expected, yamlToString(updated))
+}
+
+func TestWrite_new_map_array_deep(t *testing.T) {
+	var data = parseData(`
+b:
+  c: 2
+`)
+
+	updated := writeMap(data, []string{"b", "d", "0"}, "4")
+	assertResult(t, "4", readMap(updated, "b", []string{"d", "0"}))
+}
+
 func TestWrite_with_no_tail(t *testing.T) {
 	var data = parseData(`
 b:
   c: 2
 `)
-	write(data, "b", []string{}, "4")
+	updated := writeMap(data, []string{"b"}, "4")
 
-	b := entryInSlice(data, "b").Value
+	b := entryInSlice(updated, "b").Value
 	assertResult(t, "4", fmt.Sprintf("%v", b))
 }

yaml.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/mikefarah/yaml/Godeps/_workspace/src/github.com/spf13/cobra"
 	"github.com/mikefarah/yaml/Godeps/_workspace/src/gopkg.in/yaml.v2"
+	"github.com/op/go-logging"
 	"io/ioutil"
 	"os"
 	"strconv"
@@ -15,16 +16,26 @@ var writeInplace = false
 var writeScript = ""
 var inputJSON = false
 var outputToJSON = false
+var verbose = false
+var log = logging.MustGetLogger("yaml")
+var format = logging.MustStringFormatter(
+	`%{color}%{time:15:04:05} %{shortfunc} [%{level:.4s}]%{color:reset} %{message}`,
+)
+var backend = logging.AddModuleLevel(
+	logging.NewBackendFormatter(logging.NewLogBackend(os.Stderr, "", 0), format))
 
 func main() {
+	backend.SetLevel(logging.ERROR, "")
+	logging.SetBackend(backend)
+
 	var cmdRead = createReadCmd()
 	var cmdWrite = createWriteCmd()
 
 	var rootCmd = &cobra.Command{Use: "yaml"}
 	rootCmd.PersistentFlags().BoolVarP(&trimOutput, "trim", "t", true, "trim yaml output")
 	rootCmd.PersistentFlags().BoolVarP(&outputToJSON, "tojson", "j", false, "output as json")
 	rootCmd.PersistentFlags().BoolVarP(&inputJSON, "fromjson", "J", false, "input as json")
-
+	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "verbose mode")
 	rootCmd.AddCommand(cmdRead, cmdWrite)
 	rootCmd.Execute()
 }
@@ -77,6 +88,9 @@ a.b.e:
 }
 
 func readProperty(cmd *cobra.Command, args []string) {
+	if verbose {
+		backend.SetLevel(logging.DEBUG, "")
+	}
 	print(read(args))
 }
 
@@ -95,6 +109,9 @@ func read(args []string) interface{} {
 }
 
 func writeProperty(cmd *cobra.Command, args []string) {
+	if verbose {
+		backend.SetLevel(logging.DEBUG, "")
+	}
 	updatedData := updateYaml(args)
 	if writeInplace {
 		ioutil.WriteFile(args[0], []byte(yamlToString(updatedData)), 0644)
@@ -119,7 +136,7 @@ func updateYaml(args []string) interface{} {
 
 	for path, value := range writeCommands {
 		var paths = parsePath(path)
-		write(parsedData, paths[0], paths[1:len(paths)], value)
+		parsedData = writeMap(parsedData, paths, value)
 	}
 
 	return parsedData