Store the replace and insert annotations' "via" kwarg as callable

carvel-dev · Sep 16, 2022 · d413e22 · d413e22
1 parent 486de37
commit d413e22
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 50 deletions.
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/insert_annotation.go b/pkg/yttlibrary/overlay/insert_annotation.go
@@ -22,7 +22,7 @@ type InsertAnnotation struct {
 	newItem template.EvaluationNode
 	before  bool
 	after   bool
-	via     *starlark.Value
+	via     *starlark.Callable
 	thread  *starlark.Thread
 }
 
@@ -64,7 +64,13 @@ func NewInsertAnnotation(newItem template.EvaluationNode, thread *starlark.Threa
 			annotation.after = resultBool
 
 		case InsertAnnotationKwargVia:
-			annotation.via = &kwarg[1]
+			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(
@@ -86,29 +92,22 @@ func (a InsertAnnotation) Value(existingNode template.EvaluationNode) (interface
 		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(),
-			yamltemplate.NewGoValueWithYAML(newNode.GetValues()[0]).AsStarlarkValue(),
-		}
-
-		result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
-		if err != nil {
-			return nil, err
-		}
+	var existingVal interface{}
+	if existingNode != nil {
+		existingVal = existingNode.DeepCopyAsInterface().(template.EvaluationNode).GetValues()[0]
+	} else {
+		existingVal = nil
+	}
 
-		return tplcore.NewStarlarkValue(result).AsGoValue()
+	viaArgs := starlark.Tuple{
+		yamltemplate.NewGoValueWithYAML(existingVal).AsStarlarkValue(),
+		yamltemplate.NewGoValueWithYAML(newNode.GetValues()[0]).AsStarlarkValue(),
+	}
 
-	default:
-		return nil, fmt.Errorf("Expected '%s' annotation keyword argument 'via'"+
-			" to be function, but was %T", AnnotationInsert, typedVal)
+	result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
+	if err != nil {
+		return nil, err
 	}
+
+	return tplcore.NewStarlarkValue(result).AsGoValue()
 }
diff --git a/pkg/yttlibrary/overlay/replace_annotation.go b/pkg/yttlibrary/overlay/replace_annotation.go
@@ -20,7 +20,7 @@ const (
 type ReplaceAnnotation struct {
 	newNode template.EvaluationNode
 	thread  *starlark.Thread
-	via     *starlark.Value
+	via     *starlark.Callable
 	orAdd   bool
 }
 
@@ -36,7 +36,13 @@ func NewReplaceAnnotation(newNode template.EvaluationNode, thread *starlark.Thre
 
 		switch kwargName {
 		case ReplaceAnnotationKwargVia:
-			annotation.via = &kwarg[1]
+			via, ok := kwarg[1].(starlark.Callable)
+			if !ok {
+				return ReplaceAnnotation{}, fmt.Errorf(
+					"Expected '%s' annotation keyword argument '%s' to be function, "+
+						"but was %T", AnnotationReplace, kwargName, kwarg[1])
+			}
+			annotation.via = &via
 
 		case ReplaceAnnotationKwargOrAdd:
 			resultBool, err := tplcore.NewStarlarkValue(kwarg[1]).AsBool()
@@ -63,33 +69,26 @@ func (a ReplaceAnnotation) Value(existingNode template.EvaluationNode) (interfac
 		return newNode.GetValues()[0], nil
 	}
 
-	switch typedVal := (*a.via).(type) {
-	case starlark.Callable:
-		var existingVal interface{}
-		if existingNode != nil {
-			// Make sure original nodes are not affected in any way
-			existingVal = existingNode.DeepCopyAsInterface().(template.EvaluationNode).GetValues()[0]
-		} else {
-			existingVal = nil
-		}
-
-		viaArgs := starlark.Tuple{
-			yamltemplate.NewGoValueWithYAML(existingVal).AsStarlarkValue(),
-			yamltemplate.NewGoValueWithYAML(newNode.GetValues()[0]).AsStarlarkValue(),
-		}
-
-		// TODO check thread correctness
-		result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
-		if err != nil {
-			return nil, err
-		}
+	var existingVal interface{}
+	if existingNode != nil {
+		// Make sure original nodes are not affected in any way
+		existingVal = existingNode.DeepCopyAsInterface().(template.EvaluationNode).GetValues()[0]
+	} else {
+		existingVal = nil
+	}
 
-		return tplcore.NewStarlarkValue(result).AsGoValue()
+	viaArgs := starlark.Tuple{
+		yamltemplate.NewGoValueWithYAML(existingVal).AsStarlarkValue(),
+		yamltemplate.NewGoValueWithYAML(newNode.GetValues()[0]).AsStarlarkValue(),
+	}
 
-	default:
-		return nil, fmt.Errorf("Expected '%s' annotation keyword argument 'via'"+
-			" to be function, but was %T", AnnotationReplace, typedVal)
+	// TODO check thread correctness
+	result, err := starlark.Call(a.thread, *a.via, viaArgs, []starlark.Tuple{})
+	if err != nil {
+		return nil, err
 	}
+
+	return tplcore.NewStarlarkValue(result).AsGoValue()
 }
 
 func (a ReplaceAnnotation) OrAdd() bool { return a.orAdd }