Optionally overlay/insert via function

The insert overlay action has an optional 'via' kwarg, which receives the left node and produces the right node, similar to the replace overlay action.
carvel-dev · Sep 11, 2022 · cfc93af · cfc93af
1 parent a787f84
commit cfc93af
Show file tree

Hide file tree

Showing 5 changed files with 171 additions and 11 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,40 @@
+#@ load("@ytt:overlay", "overlay")
+
+---
+#@ def test1_left():
+- item: 1
+- item: 2
+- item: 3
+#@ end
+
+---
+#@ def test1_right():
+#@overlay/match by=lambda i, l, r: "item" in l, expects=3
+#@overlay/insert before=True, via=lambda l: {"from_item": l["item"], "before": True}
+- {}
+#@overlay/match by=lambda i, l, r: "item" in l, expects=3
+#@overlay/insert after=True, via=lambda l: {"from_item": l["item"], "after": True}
+- {}
+#@ end
+
+---
+test1: #@ overlay.apply(test1_left(), test1_right())
+
++++
+
+test1:
+- from_item: 1
+  before: true
+- item: 1
+- from_item: 1
+  after: true
+- from_item: 2
+  before: true
+- item: 2
+- from_item: 2
+  after: true
+- from_item: 3
+  before: true
+- item: 3
+- from_item: 3
+  after: true
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,48 @@
+#@ load("@ytt:template", "template")
+#@ load("@ytt:overlay", "overlay")
+
+#@ def test1_left():
+---
+item: 1
+---
+item: 2
+---
+item: 3
+#@ end
+
+#@ def test1_right():
+#@overlay/match by=lambda i, l, r: "item" in l, expects=3
+#@overlay/insert before=True, via=lambda l: {"from_item": l["item"], "before": True}
+---
+#@overlay/match by=lambda i, l, r: "item" in l, expects=3
+#@overlay/insert after=True, via=lambda l: {"from_item": l["item"], "after": True}
+---
+#@ end
+
+--- #@ template.replace(overlay.apply(test1_left(), test1_right()))
+
++++
+
+from_item: 1
+before: true
+---
+item: 1
+---
+from_item: 1
+after: true
+---
+from_item: 2
+before: true
+---
+item: 2
+---
+from_item: 2
+after: true
+---
+from_item: 3
+before: true
+---
+item: 3
+---
+from_item: 3
+after: true
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.Value
+	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,23 @@ 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:
+			annotation.via = &kwarg[1]
+
 		default:
 			return annotation, fmt.Errorf(
 				"Unknown '%s' annotation keyword argument '%s'", AnnotationInsert, kwargName)
@@ -61,3 +77,37 @@ 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
+	}
+
+	switch typedVal := (*a.via).(type) {
+	case starlark.Callable:
+		var existingVal interface{}
+		if existingNode != nil {
+			existingVal = existingNode.DeepCopyAsInterface().(template.EvaluationNode).GetValues()[0]
+		} else {
+			existingVal = nil
+		}
+
+		viaArgs := starlark.Tuple{
+			yamltemplate.NewGoValueWithYAML(existingVal).AsStarlarkValue(),
+		}
+
+		result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
+		if err != nil {
+			return nil, err
+		}
+
+		return tplcore.NewStarlarkValue(result).AsGoValue()
+
+	default:
+		return nil, fmt.Errorf("Expected '%s' annotation keyword argument 'via'"+
+			" to be function, but was %T", AnnotationInsert, typedVal)
+	}
+}