Merge pull request #742 from mamachanko/insert-via

Optionally, `overlay/insert` via function
carvel-dev · Sep 22, 2022 · edd111f · edd111f
2 parents f5ffa14 + d413e22
commit edd111f
Show file tree

Hide file tree

Showing 8 changed files with 306 additions and 37 deletions.
diff --git a/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-array-via.tpltest b/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-array-via.tpltest
@@ -0,0 +1,72 @@
+#@ load("@ytt:overlay", "overlay")
+
+---
+#@ def test1_left():
+- item: 1
+- item: 2
+- item: 3
+#@ end
+
+#@ def is_item(indexOrKey, left, right):
+#@   return "item" in left
+#@ end
+
+#@ def insert_function(left, right):
+#@   return {"left": left, "right": right, "combined": {"item": left["item"] + right["item"]}}
+#@ end
+
+---
+#@ def test1_right():
+#@overlay/match by=is_item, expects=3
+#@overlay/insert before=True, via=insert_function
+- item: 100
+#@overlay/match by=is_item, expects=3
+#@overlay/insert after=True, via=insert_function
+- item: 200
+#@ end
+
+---
+test1: #@ overlay.apply(test1_left(), test1_right())
+
++++
+
+test1:
+- left:
+    item: 1
+  right:
+    item: 100
+  combined:
+    item: 101
+- item: 1
+- left:
+    item: 1
+  right:
+    item: 200
+  combined:
+    item: 201
+- left:
+    item: 2
+  right:
+    item: 100
+  combined:
+    item: 102
+- item: 2
+- left:
+    item: 2
+  right:
+    item: 200
+  combined:
+    item: 202
+- left:
+    item: 3
+  right:
+    item: 100
+  combined:
+    item: 103
+- item: 3
+- left:
+    item: 3
+  right:
+    item: 200
+  combined:
+    item: 203
diff --git a/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-doc-via.tpltest b/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-doc-via.tpltest
@@ -0,0 +1,82 @@
+#@ load("@ytt:template", "template")
+#@ load("@ytt:overlay", "overlay")
+
+#@ def test1_left():
+---
+item: 1
+---
+item: 2
+---
+item: 3
+#@ end
+
+#@ def is_item(indexOrKey, left, right):
+#@   return "item" in left
+#@ end
+
+#@ def insert_function(left, right):
+#@   return {"left": left, "right": right, "combined": {"item": left["item"] + right["item"]}}
+#@ end
+
+#@ def test1_right():
+#@overlay/match by=is_item, expects=3
+#@overlay/insert before=True, via=insert_function
+---
+item: 100
+#@overlay/match by=is_item, expects=3
+#@overlay/insert after=True, via=insert_function
+---
+item: 200
+#@ end
+
+--- #@ template.replace(overlay.apply(test1_left(), test1_right()))
+
++++
+
+left:
+  item: 1
+right:
+  item: 100
+combined:
+  item: 101
+---
+item: 1
+---
+left:
+  item: 1
+right:
+  item: 200
+combined:
+  item: 201
+---
+left:
+  item: 2
+right:
+  item: 100
+combined:
+  item: 102
+---
+item: 2
+---
+left:
+  item: 2
+right:
+  item: 200
+combined:
+  item: 202
+---
+left:
+  item: 3
+right:
+  item: 100
+combined:
+  item: 103
+---
+item: 3
+---
+left:
+  item: 3
+right:
+  item: 200
+combined:
+  item: 203
diff --git a/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-via-err.tpltest b/pkg/yamltemplate/filetests/ytt-library/overlay/insert-into-via-err.tpltest
@@ -0,0 +1,22 @@
+#@ load("@ytt:template", "template")
+#@ load("@ytt:overlay", "overlay")
+
+#@ def test1_left():
+---
+item: 1
+#@ end
+
+#@ def test1_right():
+#@overlay/match by=overlay.all, expects=1
+#@overlay/insert before=True, via="not a function"
+---
+#@ end
+
+--- #@ template.replace(overlay.apply(test1_left(), test1_right()))
+
++++
+
+ERR:
+- overlay.apply: Document on line stdin:12: Expected 'overlay/insert' annotation keyword argument 'via' to be function, but was starlark.String
+    in <toplevel>
+      stdin:15 | --- #@ template.replace(overlay.apply(test1_left(), test1_right()))
diff --git a/pkg/yamltemplate/filetests/ytt-library/overlay/replace-func-err.tpltest b/pkg/yamltemplate/filetests/ytt-library/overlay/replace-func-err.tpltest
@@ -0,0 +1,22 @@
+#@ load("@ytt:template", "template")
+#@ load("@ytt:overlay", "overlay")
+
+#@ def test1_left():
+---
+item: 1
+#@ end
+
+#@ def test1_right():
+#@overlay/match by=overlay.all, expects=1
+#@overlay/replace via="not a function"
+---
+#@ end
+
+--- #@ template.replace(overlay.apply(test1_left(), test1_right()))
+
++++
+
+ERR:
+- overlay.apply: Document on line stdin:12: Expected 'overlay/replace' annotation keyword argument 'via' to be function, but was starlark.String
+    in <toplevel>
+      stdin:15 | --- #@ template.replace(overlay.apply(test1_left(), test1_right()))
diff --git a/pkg/yttlibrary/overlay/array.go b/pkg/yttlibrary/overlay/array.go
@@ -159,7 +159,7 @@ func (o Op) insertArrayItem(
 		return err
 	}
 
-	insertAnn, err := NewInsertAnnotation(newItem)
+	insertAnn, err := NewInsertAnnotation(newItem, o.Thread)
 	if err != nil {
 		return err
 	}
@@ -171,12 +171,23 @@ func (o Op) insertArrayItem(
 		for _, leftIdx := range leftIdxs {
 			if i == leftIdx {
 				matched = true
+
+				newVal, err := insertAnn.Value(leftItem)
+				if err != nil {
+					return err
+				}
+				insertItem := newItem.DeepCopy()
+				err = insertItem.SetValue(newVal)
+				if err != nil {
+					return err
+				}
+
 				if insertAnn.IsBefore() {
-					updatedItems = append(updatedItems, newItem.DeepCopy())
+					updatedItems = append(updatedItems, insertItem)
 				}
 				updatedItems = append(updatedItems, leftItem)
 				if insertAnn.IsAfter() {
-					updatedItems = append(updatedItems, newItem.DeepCopy())
+					updatedItems = append(updatedItems, insertItem)
 				}
 				break
 			}

diff --git a/pkg/yttlibrary/overlay/document.go b/pkg/yttlibrary/overlay/document.go
@@ -156,7 +156,7 @@ func (o Op) insertDocument(
 		return err
 	}
 
-	insertAnn, err := NewInsertAnnotation(newDoc)
+	insertAnn, err := NewInsertAnnotation(newDoc, o.Thread)
 	if err != nil {
 		return err
 	}
@@ -169,12 +169,23 @@ func (o Op) insertDocument(
 			for _, leftIdx := range leftIdxs {
 				if leftIdx[0] == i && leftIdx[1] == j {
 					matched = true
+
+					newVal, err := insertAnn.Value(leftItem)
+					if err != nil {
+						return err
+					}
+					insertDoc := newDoc.DeepCopy()
+					err = insertDoc.SetValue(newVal)
+					if err != nil {
+						return err
+					}
+
 					if insertAnn.IsBefore() {
-						updatedDocs = append(updatedDocs, newDoc.DeepCopy())
+						updatedDocs = append(updatedDocs, insertDoc)
 					}
 					updatedDocs = append(updatedDocs, leftItem)
 					if insertAnn.IsAfter() {
-						updatedDocs = append(updatedDocs, newDoc.DeepCopy())
+						updatedDocs = append(updatedDocs, insertDoc)
 					}
 					break
 				}

diff --git a/pkg/yttlibrary/overlay/insert_annotation.go b/pkg/yttlibrary/overlay/insert_annotation.go
@@ -5,20 +5,33 @@ package overlay
 
 import (
 	"fmt"
-
 	"github.com/k14s/starlark-go/starlark"
 	"github.com/vmware-tanzu/carvel-ytt/pkg/template"
 	tplcore "github.com/vmware-tanzu/carvel-ytt/pkg/template/core"
+	"github.com/vmware-tanzu/carvel-ytt/pkg/yamltemplate"
+)
+
+// Kwargs of overlay/insert
+const (
+	InsertAnnotationKwargBefore string = "before"
+	InsertAnnotationKwargAfter  string = "after"
+	InsertAnnotationKwargVia    string = "via"
 )
 
 type InsertAnnotation struct {
 	newItem template.EvaluationNode
 	before  bool
 	after   bool
+	via     *starlark.Callable
+	thread  *starlark.Thread
 }
 
-func NewInsertAnnotation(newItem template.EvaluationNode) (InsertAnnotation, error) {
-	annotation := InsertAnnotation{newItem: newItem}
+// NewInsertAnnotation returns a new InsertAnnotation for the given node and with the given Starlark thread
+func NewInsertAnnotation(newItem template.EvaluationNode, thread *starlark.Thread) (InsertAnnotation, error) {
+	annotation := InsertAnnotation{
+		newItem: newItem,
+		thread:  thread,
+	}
 	anns := template.NewAnnotations(newItem)
 
 	if !anns.Has(AnnotationInsert) {
@@ -36,20 +49,29 @@ func NewInsertAnnotation(newItem template.EvaluationNode) (InsertAnnotation, err
 		kwargName := string(kwarg[0].(starlark.String))
 
 		switch kwargName {
-		case "before":
+		case InsertAnnotationKwargBefore:
 			resultBool, err := tplcore.NewStarlarkValue(kwarg[1]).AsBool()
 			if err != nil {
 				return InsertAnnotation{}, err
 			}
 			annotation.before = resultBool
 
-		case "after":
+		case InsertAnnotationKwargAfter:
 			resultBool, err := tplcore.NewStarlarkValue(kwarg[1]).AsBool()
 			if err != nil {
 				return InsertAnnotation{}, err
 			}
 			annotation.after = resultBool
 
+		case InsertAnnotationKwargVia:
+			via, ok := kwarg[1].(starlark.Callable)
+			if !ok {
+				return InsertAnnotation{}, fmt.Errorf(
+					"Expected '%s' annotation keyword argument '%s' to be function, "+
+						"but was %T", AnnotationInsert, kwargName, kwarg[1])
+			}
+			annotation.via = &via
+
 		default:
 			return annotation, fmt.Errorf(
 				"Unknown '%s' annotation keyword argument '%s'", AnnotationInsert, kwargName)
@@ -61,3 +83,31 @@ func NewInsertAnnotation(newItem template.EvaluationNode) (InsertAnnotation, err
 
 func (a InsertAnnotation) IsBefore() bool { return a.before }
 func (a InsertAnnotation) IsAfter() bool  { return a.after }
+
+// Value returns the new value for the given, existing node. If `via` is not provided, the value of the existing
+// node is returned, otherwise the result of `via`.
+func (a InsertAnnotation) Value(existingNode template.EvaluationNode) (interface{}, error) {
+	newNode := a.newItem.DeepCopyAsInterface().(template.EvaluationNode)
+	if a.via == nil {
+		return newNode.GetValues()[0], nil
+	}
+
+	var existingVal interface{}
+	if existingNode != nil {
+		existingVal = existingNode.DeepCopyAsInterface().(template.EvaluationNode).GetValues()[0]
+	} else {
+		existingVal = nil
+	}
+
+	viaArgs := starlark.Tuple{
+		yamltemplate.NewGoValueWithYAML(existingVal).AsStarlarkValue(),
+		yamltemplate.NewGoValueWithYAML(newNode.GetValues()[0]).AsStarlarkValue(),
+	}
+
+	result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
+	if err != nil {
+		return nil, err
+	}
+
+	return tplcore.NewStarlarkValue(result).AsGoValue()
+}