From eb5c65715ddaf1b1524d303296950cc7c9f16b19 Mon Sep 17 00:00:00 2001
From: John Brunton <john_brunton@hotmail.co.uk>
Date: Thu, 25 Jun 2020 12:12:02 +0100
Subject: [PATCH] add suggestions template

---
 command.go      | 31 +++++++++++++++++++++++++++----
 command_test.go | 19 +++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/command.go b/command.go
index 5f1caccc..8596cff3 100644
--- a/command.go
+++ b/command.go
@@ -201,6 +201,8 @@ type Command struct {
 	helpCommand *Command
 	// versionTemplate is the version template defined by user.
 	versionTemplate string
+	// suggestionsTemplate is the suggestions template defined by user.
+	suggestionsTemplate string
 
 	// inReader is a reader defined by the user that replaces stdin
 	inReader io.Reader
@@ -284,6 +286,11 @@ func (c *Command) SetVersionTemplate(s string) {
 	c.versionTemplate = s
 }
 
+// SetSuggestionsTemplate sets suggestions template. Can be defined by Application.
+func (c *Command) SetSuggestionsTemplate(s string) {
+	c.suggestionsTemplate = s
+}
+
 // SetGlobalNormalizationFunc sets a normalization function to all flag sets and also to child commands.
 // The user should not have a cyclic dependency on commands.
 func (c *Command) SetGlobalNormalizationFunc(n func(f *flag.FlagSet, name string) flag.NormalizedName) {
@@ -527,6 +534,22 @@ func (c *Command) VersionTemplate() string {
 `
 }
 
+// SuggestionsTemplate return suggestions template for the command.
+func (c *Command) SuggestionsTemplate() string {
+	if c.suggestionsTemplate != "" {
+		return c.suggestionsTemplate
+	}
+
+	if c.HasParent() {
+		return c.parent.SuggestionsTemplate()
+	}
+	return `
+
+Did you mean this?
+{{range .}}{{print "\t" .}}{{end}}
+`
+}
+
 func hasNoOptDefVal(name string, fs *flag.FlagSet) bool {
 	flag := fs.Lookup(name)
 	if flag == nil {
@@ -639,10 +662,10 @@ func (c *Command) findSuggestions(arg string) string {
 	}
 	suggestionsString := ""
 	if suggestions := c.SuggestionsFor(arg); len(suggestions) > 0 {
-		suggestionsString += "\n\nDid you mean this?\n"
-		for _, s := range suggestions {
-			suggestionsString += fmt.Sprintf("\t%v\n", s)
-		}
+		var suggestionsBuffer bytes.Buffer
+		suggestions := c.SuggestionsFor(arg)
+		tmpl(&suggestionsBuffer, c.SuggestionsTemplate(), suggestions)
+		suggestionsString += suggestionsBuffer.String()
 	}
 	return suggestionsString
 }
diff --git a/command_test.go b/command_test.go
index 16cc41b4..0658dcac 100644
--- a/command_test.go
+++ b/command_test.go
@@ -1202,6 +1202,25 @@ func TestSuggestions(t *testing.T) {
 	}
 }
 
+func TestSuggestionsTemplate(t *testing.T) {
+	rootCmd := &Command{Use: "root", Run: emptyRun}
+	timesCmd := &Command{
+		Use:        "times",
+		SuggestFor: []string{"counts"},
+		Run:        emptyRun,
+	}
+	rootCmd.AddCommand(timesCmd)
+	rootCmd.SetSuggestionsTemplate(`
+customized suggestions: {{range .}}{{.}}{{end}}`)
+
+	output, _ := executeCommand(rootCmd, "time")
+	expected := "Error: unknown command \"time\" for \"root\"\ncustomized suggestions: times\nRun 'root --help' for usage.\n"
+
+	if output != expected {
+		t.Errorf("Unexpected response.\nExpected:\n %q\nGot:\n %q\n", expected, output)
+	}
+}
+
 func TestRemoveCommand(t *testing.T) {
 	rootCmd := &Command{Use: "root", Args: NoArgs, Run: emptyRun}
 	childCmd := &Command{Use: "child", Run: emptyRun}