From 5b121bc9fbffd5ba668b5aa101790d4b27917d84 Mon Sep 17 00:00:00 2001
From: akutz <sakutz@gmail.com>
Date: Mon, 31 Aug 2015 22:36:55 -0500
Subject: [PATCH] Template Function Injection

This patch enables developers to add one to many template functions that
can be used by custom Usage and Help templates. Here is an example that
is included in the file cobra_test.go as the test function named
TestAddTemplateFunctions:

    AddTemplateFunc("t", func() bool { return true })
    AddTemplateFuncs(template.FuncMap{
        "f": func() bool { return false },
        "h": func() string { return "Hello," },
        "w": func() string { return "world." }})

    const usage = "Hello, world."

    c := &Command{}
    c.SetUsageTemplate(`{{if t}}{{h}}{{end}}{{if f}}{{h}}{{end}} {{w}}`)

    if us := c.UsageString(); us != usage {
        t.Errorf("c.UsageString() != \"%s\", is \"%s\"", usage, us)
    }

In the above example four functions are added to the template function
map used when the Usage and Help text is generated from the templates
that enable custom logic as well as data injection during template
execution.
---
 cobra.go      | 28 ++++++++++++++++++++++------
 cobra_test.go | 18 ++++++++++++++++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/cobra.go b/cobra.go
index 78b92b0..1d66a70 100644
--- a/cobra.go
+++ b/cobra.go
@@ -25,6 +25,13 @@ import (
 	"text/template"
 )
 
+var templateFuncs template.FuncMap = template.FuncMap{
+	"trim": strings.TrimSpace,
+	"rpad": rpad,
+	"gt":   Gt,
+	"eq":   Eq,
+}
+
 var initializers []func()
 
 // automatic prefix matching can be a dangerous thing to automatically enable in CLI tools.
@@ -39,6 +46,20 @@ var MousetrapHelpText string = `This is a command line tool
 You need to open cmd.exe and run it from there.
 `
 
+//AddTemplateFunc adds a template function that's available to Usage and Help
+//template generation.
+func AddTemplateFunc(name string, tmplFunc interface{}) {
+	templateFuncs[name] = tmplFunc
+}
+
+//AddTemplateFuncs adds multiple template functions availalble to Usage and
+//Help template generation.
+func AddTemplateFuncs(tmplFuncs template.FuncMap) {
+	for k, v := range tmplFuncs {
+		templateFuncs[k] = v
+	}
+}
+
 //OnInitialize takes a series of func() arguments and appends them to a slice of func().
 func OnInitialize(y ...func()) {
 	for _, x := range y {
@@ -101,12 +122,7 @@ func rpad(s string, padding int) string {
 // tmpl executes the given template text on data, writing the result to w.
 func tmpl(w io.Writer, text string, data interface{}) error {
 	t := template.New("top")
-	t.Funcs(template.FuncMap{
-		"trim": strings.TrimSpace,
-		"rpad": rpad,
-		"gt":   Gt,
-		"eq":   Eq,
-	})
+	t.Funcs(templateFuncs)
 	template.Must(t.Parse(text))
 	return t.Execute(w, data)
 }
diff --git a/cobra_test.go b/cobra_test.go
index 4fc3b88..3aed7dd 100644
--- a/cobra_test.go
+++ b/cobra_test.go
@@ -8,6 +8,7 @@ import (
 	"runtime"
 	"strings"
 	"testing"
+	"text/template"
 
 	"github.com/spf13/pflag"
 )
@@ -971,3 +972,20 @@ func TestFlagOnPflagCommandLine(t *testing.T) {
 
 	checkResultContains(t, r, flagName)
 }
+
+func TestAddTemplateFunctions(t *testing.T) {
+	AddTemplateFunc("t", func() bool { return true })
+	AddTemplateFuncs(template.FuncMap{
+		"f": func() bool { return false }, 
+		"h": func() string { return "Hello," }, 
+		"w": func() string { return "world." }})
+
+	const usage = "Hello, world."
+	
+	c := &Command{}
+	c.SetUsageTemplate(`{{if t}}{{h}}{{end}}{{if f}}{{h}}{{end}} {{w}}`)
+	
+	if us := c.UsageString(); us != usage {
+		t.Errorf("c.UsageString() != \"%s\", is \"%s\"", usage, us)
+	}
+}