Add method to insert a value at a given JSONPointer location (#3)

Add method to insert a value at a given JSONPointer location

yptr.go

@@ -4,7 +4,9 @@
 package yptr
 
 import (
+	"errors"
 	"fmt"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -78,6 +80,120 @@ func find(root *yaml.Node, toks []string) ([]*yaml.Node, error) {
 	return res, nil
 }
 
+// Insert inserts a value at the location pointed by the JSONPointer ptr in the yaml tree rooted at root.
+// If any nodes along the way do not exist, they are created such that a subsequent call to Find would find
+// the value at that location.
+//
+// Note that Insert does not replace existing values. If the location already exists in the yaml tree, Insert will
+// attempt to append the value to the existing node there. If this isn't possible, an error is returned.
+//
+// Also note that '-' is only treated as a special character if the currently referenced value is an existing array.
+// It cannot be used to create a new empty array at the current location.
+func Insert(root *yaml.Node, ptr string, value yaml.Node) error {
+	toks, err := jsonPointerToTokens(ptr)
+	if err != nil {
+		return err
+	}
+
+	// skip document nodes
+	if value.Kind == yaml.DocumentNode {
+		value = *value.Content[0]
+	}
+	if root.Kind == yaml.DocumentNode {
+		root = root.Content[0]
+	}
+
+	return insert(root, toks, value)
+}
+
+func insert(root *yaml.Node, toks []string, value yaml.Node) error {
+	if len(toks) == 0 {
+		if root.Kind == yaml.MappingNode {
+			if value.Kind == yaml.MappingNode {
+				root.Content = append(root.Content, value.Content...)
+				return nil
+			}
+			if len(root.Content) == 0 {
+				*root = value
+				return nil
+			}
+		}
+		return fmt.Errorf("cannot insert node type %v (%v) in node type %v (%v)", value.Kind, value.Tag, root.Kind, root.Tag)
+	}
+
+	switch root.Kind {
+	case yaml.SequenceNode:
+		return sequenceInsert(root, toks, value)
+	case yaml.MappingNode:
+		return mapInsert(root, toks, value)
+	default:
+		return fmt.Errorf("unhandled node type: %v (%v)", root.Kind, root.Tag)
+	}
+}
+
+func sequenceInsert(root *yaml.Node, toks []string, value yaml.Node) error {
+	// try to find the token in the node
+	next, err := match(root, toks[0])
+	if err != nil {
+		return err
+	}
+	if len(toks) == 1 {
+		return sequenceInsertAt(root, toks[0], value)
+	}
+	if next[0].Kind != yaml.ScalarNode {
+		return insert(next[0], toks[1:], value)
+	}
+	// insert an empty map and try again
+	n := yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
+	k := yaml.Node{Kind: yaml.ScalarNode, Value: toks[1], Tag: "!!str"}
+	v := yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
+	n.Content = append(n.Content, &k, &v)
+
+	err = sequenceInsertAt(root, toks[0], n)
+	if err != nil {
+		return err
+	}
+	return insert(&n, toks[1:], value)
+}
+
+// helper function for inserting a value at a specific index in an array
+func sequenceInsertAt(root *yaml.Node, tok string, n yaml.Node) error {
+	if tok == "-" {
+		root.Content = append(root.Content, &n)
+	} else {
+		i, err := strconv.Atoi(tok)
+		if err != nil {
+			return err
+		}
+		if i < 0 || i >= len(root.Content) {
+			return fmt.Errorf("out of bounds")
+		}
+		root.Content = slices.Insert(root.Content, i, &n)
+	}
+	return nil
+}
+
+func mapInsert(root *yaml.Node, toks []string, value yaml.Node) error {
+	// try to find the token in the node
+	next, err := match(root, toks[0])
+	if err != nil && !errors.Is(err, ErrNotFound) {
+		return err
+	}
+
+	if errors.Is(err, ErrNotFound) {
+		k := yaml.Node{Kind: yaml.ScalarNode, Value: toks[0]}
+		if len(toks) == 1 {
+			root.Content = append(root.Content, &k, &value)
+			return nil
+		}
+		// insert an empty map and try again
+		v := yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
+		root.Content = append(root.Content, &k, &v)
+		return insert(&v, toks[1:], value)
+	}
+	return insert(next[0], toks[1:], value)
+}
+
 // match matches a JSONPointer token against a yaml Node.
 //
 // If root is a map, it performs a field lookup using tok as field name,
@@ -111,6 +227,10 @@ func match(root *yaml.Node, tok string) ([]*yaml.Node, error) {
 			}
 			return filter(c, treeSubsetPred(&mtree))
 		default:
+			if tok == "-" {
+				// dummy leaf node
+				return []*yaml.Node{{Kind: yaml.ScalarNode}}, nil
+			}
 			i, err := strconv.Atoi(tok)
 			if err != nil {
 				return nil, err

yptr_test.go

@@ -6,12 +6,107 @@ package yptr_test
 import (
 	"errors"
 	"fmt"
+	"strings"
 	"testing"
 
 	yptr "github.com/zillow/go-yaml-jsonpointer"
 	"github.com/zillow/go-yaml/v3"
 )
 
+func ExampleInsert() {
+	src := `
+d:
+  - e
+  - f:
+      g: x
+  - h: y
+  - - i
+    - j
+`
+	arr1 := `[1, 2, 3]`
+	map1 := `q: xyz`
+	s1 := `x`
+
+	var n, a, m, x yaml.Node
+	yaml.Unmarshal([]byte(src), &n)
+	yaml.Unmarshal([]byte(arr1), &a)
+	yaml.Unmarshal([]byte(map1), &m)
+	yaml.Unmarshal([]byte(s1), &x)
+
+	_ = yptr.Insert(&n, `/f/d`, a)
+	_ = yptr.Insert(&n, ``, m)
+
+	_ = yptr.Insert(&n, `/d/1`, m)
+	_ = yptr.Insert(&n, `/d/-/c`, x)
+	_ = yptr.Insert(&n, `/d/2/f`, m)
+	_ = yptr.Insert(&n, `/d/3/f`, a)
+	_ = yptr.Insert(&n, `/d/4/-`, x)
+
+	out, err := yaml.Marshal(n.Content[0])
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(string(out))
+	/* Output:
+d:
+    - e
+    - q: xyz
+    - f:
+        g: x
+        q: xyz
+    - h: y
+      f: [1, 2, 3]
+    - - i
+      - j
+      - x
+    - c: x
+f:
+    d: [1, 2, 3]
+q: xyz
+*/
+}
+
+func TestInsertErrors(t *testing.T) {
+	src := `
+a:
+  b:
+    c: 42
+d:
+- e
+- f
+`
+	s1 := `x`
+	var n, x yaml.Node
+	yaml.Unmarshal([]byte(src), &n)
+	yaml.Unmarshal([]byte(s1), &x)
+
+	tests := []struct {
+		ptr string
+		value yaml.Node
+		err string
+	}{
+		{``, x, "cannot insert node type"},
+		{`/a/b/c`, x, "cannot insert node type"},
+		{`/d`, x, "cannot insert node type"},
+		{`/a/b/c/f`, x, "unhandled node type"},
+		{`/d/f`, x, "strconv.Atoi"},
+		{`/d/5`, x, "out of bounds"},
+	}
+
+	for i, tc := range tests {
+		t.Run(fmt.Sprint(i), func(t *testing.T) {
+			err := yptr.Insert(&n, tc.ptr, tc.value)
+			if err == nil {
+				t.Fatal("expecting error")
+			}
+			if !strings.HasPrefix(err.Error(), tc.err) {
+				t.Fatalf("expecting error %q, got %q", tc.err, err)
+			}
+		})
+	}
+}
+
 func ExampleFind() {
 	src := `
 a: