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) + } +}