From 6291a093ca6da07e92fd405ebe9446fdcba33f6e Mon Sep 17 00:00:00 2001
From: JulesDT <jules.denardou@datadoghq.com>
Date: Mon, 21 Sep 2020 19:15:06 -0400
Subject: [PATCH] add colors to CLI commands

---
 command.go      | 37 ++++++++++++++++++++++++++++++++++++-
 command_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/command.go b/command.go
index 27d39d54..e6ba8b86 100644
--- a/command.go
+++ b/command.go
@@ -31,6 +31,30 @@ import (
 // FParseErrWhitelist configures Flag parse errors to be ignored
 type FParseErrWhitelist flag.ParseErrorsWhitelist
 
+// TerminalColor is a type used for the names of the different
+// colors in the terminal
+type TerminalColor int
+
+// Colors represents the different colors one can use in the terminal
+const (
+	ColorBlack TerminalColor = iota + 30
+	ColorRed
+	ColorGreen
+	ColorYellow
+	ColorBlue
+	ColorMagenta
+	ColorCyan
+	ColorLightGray
+	ColorDarkGray
+	ColorLightRed
+	ColorLightGreen
+	ColorLightYellow
+	ColorLightBlue
+	ColorLightMagenta
+	ColorLightCyan
+	ColorWhite
+)
+
 // Command is just that, a command for your application.
 // E.g.  'go run ...' - 'run' is the command. Cobra requires
 // you to define the usage and description as part of your command
@@ -47,6 +71,9 @@ type Command struct {
 	// Example: add [-F file | -D dir]... [-f format] profile
 	Use string
 
+	// Color represents the color to use to print the command in the terminal
+	Color TerminalColor
+
 	// Aliases is an array of aliases that can be used instead of the first word in Use.
 	Aliases []string
 
@@ -493,7 +520,7 @@ Examples:
 {{.Example}}{{end}}{{if .HasAvailableSubCommands}}
 
 Available Commands:{{range .Commands}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
-  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}
+  {{rpad .ColoredName .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}
 
 Flags:
 {{.LocalFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .HasAvailableInheritedFlags}}
@@ -1293,6 +1320,14 @@ func (c *Command) Name() string {
 	return name
 }
 
+// ColoredName returns the command's Name in the correct color if specified
+func (c *Command) ColoredName() string {
+	if c.Color != 0 {
+		return fmt.Sprintf("\033[%dm%s\033[0m", c.Color, c.Name())
+	}
+	return c.Name()
+}
+
 // HasAlias determines if a given string is an alias of the command.
 func (c *Command) HasAlias(s string) bool {
 	for _, a := range c.Aliases {
diff --git a/command_test.go b/command_test.go
index 16cc41b4..a521b8b8 100644
--- a/command_test.go
+++ b/command_test.go
@@ -1989,3 +1989,30 @@ func TestFParseErrWhitelistSiblingCommand(t *testing.T) {
 	}
 	checkStringContains(t, output, "unknown flag: --unknown")
 }
+
+func TestColoredName(t *testing.T) {
+	c := &Command{
+		Use: "cmd",
+	}
+	if c.Name() != "cmd" {
+		t.Error("Unexpected name with simple Command")
+	}
+	// If no color is specified, the ColoredName should equal the Name
+	if c.Name() != c.ColoredName() {
+		t.Error("Name and ColoredName should give the same result")
+	}
+	c = &Command{
+		Use:   "cmd",
+		Color: ColorRed,
+	}
+	if c.Name() != "cmd" {
+		t.Errorf("Unexpected name with Colored Command: %s\n", c.Name())
+	}
+	// If a color is specified, the ColoredName and the Name should be different
+	if c.Name() == c.ColoredName() {
+		t.Error("Name and ColoredName should not give the same result")
+	}
+	if c.ColoredName() != "\033[31m"+c.Name()+"\033[0m" {
+		t.Error("ColoredName should only add color to the name")
+	}
+}