From aea0de0dab19107da975c3a27696c9e8b1cb8e74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Antti=20Kivim=C3=A4ki?= <antti.kivimaki@futurice.com>
Date: Wed, 27 Nov 2024 11:25:46 +0200
Subject: [PATCH] feat: add include template function

---
 docs/site/docs/usage.mdx  |  2 +-
 pkg/engine/engine.go      |  2 +-
 pkg/engine/engine_test.go |  5 +++--
 pkg/engine/funcs.go       | 25 +++++++++++++++++++++++--
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/docs/site/docs/usage.mdx b/docs/site/docs/usage.mdx
index 94c3d326..0d856243 100644
--- a/docs/site/docs/usage.mdx
+++ b/docs/site/docs/usage.mdx
@@ -89,7 +89,7 @@ The rest of the files are rendered from the templates. You can edit the template
 
 ## Templating
 
-Templates are done by using [Go templates](https://pkg.go.dev/text/template). Examples of how to write these templates can be found in [this article](https://www.digitalocean.com/community/tutorials/how-to-use-templates-in-go#step-4-writing-a-template) and in the [example recipe](https://github.com/futurice/jalapeno/blob/main/examples/variable-types/templates/README.md?plain=1).
+Templates are done by using [Go templates](https://pkg.go.dev/text/template) extended with [sprig functions](https://masterminds.github.io/sprig/). Examples of how to write these templates can be found in [this article](https://www.digitalocean.com/community/tutorials/how-to-use-templates-in-go#step-4-writing-a-template) and in the [example recipe](https://github.com/futurice/jalapeno/blob/main/examples/variable-types/templates/README.md?plain=1).
 
 The following context is available on the templates:
 
diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
index 80d314a2..1091bfb1 100644
--- a/pkg/engine/engine.go
+++ b/pkg/engine/engine.go
@@ -16,7 +16,7 @@ func New() Engine {
 
 func (e Engine) Render(templates map[string][]byte, values map[string]interface{}) (map[string][]byte, error) {
 	t := template.New("gotpl")
-	t.Funcs(funcMap())
+	t.Funcs(funcMap(t))
 
 	rendered := make(map[string][]byte)
 
diff --git a/pkg/engine/engine_test.go b/pkg/engine/engine_test.go
index 5c8a6e55..01064be5 100644
--- a/pkg/engine/engine_test.go
+++ b/pkg/engine/engine_test.go
@@ -30,13 +30,14 @@ func TestRender(t *testing.T) {
 		{
 			"macros",
 			map[string][]byte{
+				"templates/main":    []byte("{{ template \"helper1\" }} {{ template \"helper2\" }} {{ include \"helper3\" . | upper }}"),
 				"templates/helper1": []byte("{{ define \"helper1\" }}ONE{{ end }}"),
-				"templates/main":    []byte("{{ template \"helper1\" }} {{ template \"helper2\" }}"),
 				"templates/helper2": []byte("{{ define \"helper2\" }}TWO{{ end }}"),
+				"templates/helper3": []byte("{{ define \"helper3\" }}three{{ end }}"),
 			},
 			map[string]interface{}{},
 			map[string][]byte{
-				"templates/main": []byte("ONE TWO"),
+				"templates/main": []byte("ONE TWO THREE"),
 			},
 		},
 	}
diff --git a/pkg/engine/funcs.go b/pkg/engine/funcs.go
index 24252ccf..53da7bff 100644
--- a/pkg/engine/funcs.go
+++ b/pkg/engine/funcs.go
@@ -1,6 +1,7 @@
 package engine
 
 import (
+	"fmt"
 	"math/big"
 	"math/rand"
 	"strings"
@@ -11,10 +12,12 @@ import (
 
 var prngs = make(map[string]*rand.Rand)
 
-func funcMap() template.FuncMap {
+func funcMap(t *template.Template) template.FuncMap {
 	f := sprig.TxtFuncMap()
+	includedNames := make(map[string]int)
 
-	// Add additional template functions here
+	// Custom template functions
+	f["include"] = includeFun(t, includedNames)
 	f["stableRandomAlphanumeric"] = stableRandomAlphanumeric
 
 	return f
@@ -23,6 +26,7 @@ func funcMap() template.FuncMap {
 const lowerAlpha = "abcdefghijlkmnopqrstuvwxyz"
 const number = "0123456789"
 const alphanumeric = lowerAlpha + number
+const recursionMaxNums = 100
 
 // Generate a pseudo-random alphanumeric string of the given length
 // such that the sequence of strings generated by successive calls
@@ -56,3 +60,20 @@ func prngForStr(str string) *rand.Rand {
 func resetRngs() {
 	prngs = make(map[string]*rand.Rand)
 }
+
+func includeFun(t *template.Template, includedNames map[string]int) func(string, interface{}) (string, error) {
+	return func(name string, data interface{}) (string, error) {
+		var buf strings.Builder
+		if v, ok := includedNames[name]; ok {
+			if v > recursionMaxNums {
+				return "", fmt.Errorf("unable to execute template: maximum recursion depth reached for template block \"%s\"", name)
+			}
+			includedNames[name]++
+		} else {
+			includedNames[name] = 1
+		}
+		err := t.ExecuteTemplate(&buf, name, data)
+		includedNames[name]--
+		return buf.String(), err
+	}
+}