Template context (#373)

Signed-off-by: David Chung <[email protected]>

Author: David Chung

pkg/template/funcs.go

@@ -3,6 +3,7 @@ package template
 import (
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"strings"
 	"time"
 
@@ -75,6 +76,35 @@ func UnixTime() interface{} {
 	return time.Now().Unix()
 }
 
+// Index returns the index of search in array.  -1 if not found or array is not iterable.  An optional true will
+// turn on strict type check while by default string representations are used to compare values.
+func Index(srch interface{}, array interface{}, strictOptional ...bool) int {
+	strict := false
+	if len(strictOptional) > 0 {
+		strict = strictOptional[0]
+	}
+	switch reflect.TypeOf(array).Kind() {
+	case reflect.Slice:
+		s := reflect.ValueOf(array)
+		for i := 0; i < s.Len(); i++ {
+			if reflect.DeepEqual(srch, s.Index(i).Interface()) {
+				return i
+			}
+			if !strict {
+				// by string value which is useful for text based compares
+				search := reflect.Indirect(reflect.ValueOf(srch)).Interface()
+				value := reflect.Indirect(s.Index(i)).Interface()
+				searchStr := fmt.Sprintf("%v", search)
+				check := fmt.Sprintf("%v", value)
+				if searchStr == check {
+					return i
+				}
+			}
+		}
+	}
+	return -1
+}
+
 // DefaultFuncs returns a list of default functions for binding in the template
 func (t *Template) DefaultFuncs() map[string]interface{} {
 	return map[string]interface{}{
@@ -102,6 +132,10 @@ func (t *Template) DefaultFuncs() map[string]interface{} {
 			return included.Render(o)
 		},
 
+		"loop": func(c int) []struct{} {
+			return make([]struct{}, c)
+		},
+
 		"var": func(name, doc string, v ...interface{}) interface{} {
 			if found, has := t.binds[name]; has {
 				return found
@@ -119,5 +153,6 @@ func (t *Template) DefaultFuncs() map[string]interface{} {
 		"lines":     SplitLines,
 		"to_json":   ToJSON,
 		"from_json": FromJSON,
+		"index":     Index,
 	}
 }

pkg/template/funcs_test.go

@@ -280,3 +280,30 @@ func TestMapEncodeDecode(t *testing.T) {
 
 	require.Equal(t, expect, actual)
 }
+
+func TestIndex(t *testing.T) {
+	require.Equal(t, -1, Index("a", []string{"x", "y", "z"}))
+	require.Equal(t, 1, Index("y", []string{"x", "y", "z"}))
+	require.Equal(t, -1, Index(25, []string{"x", "y", "z"}))
+	require.Equal(t, -1, Index(25, 26))
+	require.Equal(t, 1, Index("y", []string{"x", "y", "z"}))
+	require.Equal(t, 1, Index("y", []interface{}{"x", "y", "z"}))
+	require.Equal(t, 1, Index(1, []interface{}{0, 1, 2}))
+	require.Equal(t, 1, Index("1", []interface{}{0, 1, 2}))
+	require.Equal(t, 1, Index(1, []interface{}{0, "1", 2}))
+	require.Equal(t, -1, Index("1", []interface{}{0, 1, 2}, true))  // strict case type must match
+	require.Equal(t, 1, Index("1", []interface{}{0, "1", 2}, true)) // strict case type must match
+	require.Equal(t, -1, Index(1, []interface{}{0, "1", 2}, true))  // strict case type must match
+
+	v := "1"
+	require.Equal(t, 1, Index(&v, []interface{}{0, "1", 2}))
+	require.Equal(t, 1, Index(&v, []interface{}{0, &v, 2}, true))
+	require.Equal(t, 1, Index(&v, []interface{}{0, &v, 2}))
+
+	a := "0"
+	c := "2"
+	require.Equal(t, 1, Index("1", []*string{&a, &v, &c}))
+
+	// This doesn't work because the type information is gone and we have just an address
+	require.Equal(t, -1, Index("1", []interface{}{0, &v, 2}))
+}

pkg/template/integration_test.go

@@ -120,3 +120,41 @@ systemctl start docker
 `,
 	}
 )
+
+func TestTemplateContext(t *testing.T) {
+
+	s := `
+{{ inc }}
+
+{{ setString "hello" }}
+
+{{ setBool true }}
+
+{{ range loop 10 }}
+  {{ inc }}
+{{ end }}
+
+The count is {{count}}
+The message is {{str}}
+
+{{ dec }}
+{{ range loop 5 }}
+  {{ dec }}
+{{ end }}
+
+The count is {{count}}
+The message is {{str}}
+`
+
+	tt, err := NewTemplate("str://"+s, Options{})
+	require.NoError(t, err)
+
+	context := &context{}
+
+	_, err = tt.Render(context)
+	require.NoError(t, err)
+
+	require.Equal(t, 5, context.Count)
+	require.True(t, context.Bool)
+	require.Equal(t, 23, context.invokes) // note this is private state not accessible in template
+}

pkg/template/template.go

@@ -2,7 +2,9 @@ package template
 
 import (
 	"bytes"
+	"fmt"
 	"io"
+	"reflect"
 	"strings"
 	"sync"
 	"text/template"
@@ -11,6 +13,30 @@ import (
 	log "github.com/Sirupsen/logrus"
 )
 
+// Function contains the description of an exported template function
+type Function struct {
+
+	// Name is the function name to bind in the template
+	Name string
+
+	// Description provides help for the function
+	Description string
+
+	// Func is the reference to the actual function
+	Func interface{}
+}
+
+// Context is a marker interface for a user-defined struct that is passed into the template engine (as context)
+// and accessible in the exported template functions.  Template functions can have the signature
+// func(template.Context, arg1, arg2 ...) (string, error) and when functions like this are registered, the template
+// engine will dynamically create and export a function of the form func(arg1, arg2...) (string, error) where
+// the context instance becomes an out-of-band struct that can be mutated by functions.  This in essence allows
+// structured data as output of the template, in addition to a string from evaluating the template.
+type Context interface {
+	// Funcs returns a list of special template functions of the form func(template.Context, arg1, arg2) interface{}
+	Funcs() []Function
+}
+
 // Options contains parameters for customizing the behavior of the engine
 type Options struct {
 
@@ -87,10 +113,10 @@ func (t *Template) Validate() (*Template, error) {
 	t.lock.Lock()
 	t.parsed = nil
 	t.lock.Unlock()
-	return t, t.build()
+	return t, t.build(nil)
 }
 
-func (t *Template) build() error {
+func (t *Template) build(context Context) error {
 	t.lock.Lock()
 	defer t.lock.Unlock()
 
@@ -105,7 +131,21 @@ func (t *Template) build() error {
 	}
 
 	for k, v := range t.funcs {
-		fm[k] = v
+		if tf, err := makeTemplateFunc(context, v); err == nil {
+			fm[k] = tf
+		} else {
+			return err
+		}
+	}
+
+	if context != nil {
+		for _, f := range context.Funcs() {
+			if tf, err := makeTemplateFunc(context, f.Func); err == nil {
+				fm[f.Name] = tf
+			} else {
+				return err
+			}
+		}
 	}
 
 	parsed, err := template.New(t.url).Funcs(fm).Parse(string(t.body))
@@ -119,18 +159,75 @@ func (t *Template) build() error {
 
 // Execute is a drop-in replace of the execute method of template
 func (t *Template) Execute(output io.Writer, context interface{}) error {
-	if err := t.build(); err != nil {
+	if err := t.build(toContext(context)); err != nil {
 		return err
 	}
 	return t.parsed.Execute(output, context)
 }
 
+func toContext(in interface{}) Context {
+	var context Context
+	if in != nil {
+		if s, is := in.(Context); is {
+			context = s
+		}
+	}
+	return context
+}
+
 // Render renders the template given the context
 func (t *Template) Render(context interface{}) (string, error) {
-	if err := t.build(); err != nil {
+	if err := t.build(toContext(context)); err != nil {
 		return "", err
 	}
 	var buff bytes.Buffer
 	err := t.parsed.Execute(&buff, context)
 	return buff.String(), err
 }
+
+// converts a function of f(Context, ags...) to a regular template function
+func makeTemplateFunc(ctx Context, f interface{}) (interface{}, error) {
+
+	contextType := reflect.TypeOf((*Context)(nil)).Elem()
+
+	ff := reflect.Indirect(reflect.ValueOf(f))
+	// first we check to see if f has the special signature where the first
+	// parameter is the context parameter...
+	if ff.Kind() != reflect.Func {
+		return nil, fmt.Errorf("not a function:%v", f)
+	}
+
+	if ff.Type().In(0).AssignableTo(contextType) {
+
+		in := make([]reflect.Type, ff.Type().NumIn()-1) // exclude the context param
+		out := make([]reflect.Type, ff.Type().NumOut())
+
+		for i := 1; i < ff.Type().NumIn(); i++ {
+			in[i-1] = ff.Type().In(i)
+		}
+		variadic := false
+		if len(in) > 0 {
+			variadic = in[len(in)-1].Kind() == reflect.Slice
+		}
+		for i := 0; i < ff.Type().NumOut(); i++ {
+			out[i] = ff.Type().Out(i)
+		}
+		funcType := reflect.FuncOf(in, out, variadic)
+		funcImpl := func(in []reflect.Value) []reflect.Value {
+			if !variadic {
+				return ff.Call(append([]reflect.Value{reflect.ValueOf(ctx)}, in...))
+			}
+
+			variadicParam := in[len(in)-1]
+			last := make([]reflect.Value, variadicParam.Len())
+			for i := 0; i < variadicParam.Len(); i++ {
+				last[i] = variadicParam.Index(i)
+			}
+			return ff.Call(append(append([]reflect.Value{reflect.ValueOf(ctx)}, in[0:len(in)-1]...), last...))
+		}
+
+		newFunc := reflect.MakeFunc(funcType, funcImpl)
+		return newFunc.Interface(), nil
+	}
+	return ff.Interface(), nil
+}