Added support for -force-color to the CLI. (#10975)

Author: apollo13

.changelog/10975.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+cli: Added support for `-force-color` to the CLI to force colored output.
+```

command/commands.go

@@ -13,6 +13,9 @@ import (
 const (
 	// EnvNomadCLINoColor is an env var that toggles colored UI output.
 	EnvNomadCLINoColor = `NOMAD_CLI_NO_COLOR`
+
+	// EnvNomadCLIForceColor is an env var that forces colored UI output.
+	EnvNomadCLIForceColor = `NOMAD_CLI_FORCE_COLOR`
 )
 
 // DeprecatedCommand is a command that wraps an existing command and prints a

command/meta.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/hashicorp/nomad/api"
+	colorable "github.com/mattn/go-colorable"
 	"github.com/mitchellh/cli"
 	"github.com/mitchellh/colorstring"
 	"github.com/posener/complete"
@@ -39,6 +40,9 @@ type Meta struct {
 	// Whether to not-colorize output
 	noColor bool
 
+	// Whether to force colorized output
+	forceColor bool
+
 	// The region to send API requests
 	region string
 
@@ -70,6 +74,7 @@ func (m *Meta) FlagSet(n string, fs FlagSetFlags) *flag.FlagSet {
 		f.StringVar(&m.region, "region", "", "")
 		f.StringVar(&m.namespace, "namespace", "", "")
 		f.BoolVar(&m.noColor, "no-color", false, "")
+		f.BoolVar(&m.forceColor, "force-color", false, "")
 		f.StringVar(&m.caCert, "ca-cert", "", "")
 		f.StringVar(&m.caPath, "ca-path", "", "")
 		f.StringVar(&m.clientCert, "client-cert", "", "")
@@ -97,6 +102,7 @@ func (m *Meta) AutocompleteFlags(fs FlagSetFlags) complete.Flags {
 		"-region":          complete.PredictAnything,
 		"-namespace":       NamespacePredictor(m.Client, nil),
 		"-no-color":        complete.PredictNothing,
+		"-force-color":     complete.PredictNothing,
 		"-ca-cert":         complete.PredictFiles("*"),
 		"-ca-path":         complete.PredictDirs("*"),
 		"-client-cert":     complete.PredictFiles("*"),
@@ -155,15 +161,47 @@ func (m *Meta) allNamespaces() bool {
 
 func (m *Meta) Colorize() *colorstring.Colorize {
 	_, coloredUi := m.Ui.(*cli.ColoredUi)
-	noColor := m.noColor || !coloredUi || !terminal.IsTerminal(int(os.Stdout.Fd()))
 
 	return &colorstring.Colorize{
 		Colors:  colorstring.DefaultColors,
-		Disable: noColor,
+		Disable: !coloredUi,
 		Reset:   true,
 	}
 }
 
+func (m *Meta) SetupUi(args []string) {
+	noColor := os.Getenv(EnvNomadCLINoColor) != ""
+	forceColor := os.Getenv(EnvNomadCLIForceColor) != ""
+
+	for _, arg := range args {
+		// Check if color is set
+		if arg == "-no-color" || arg == "--no-color" {
+			noColor = true
+		} else if arg == "-force-color" || arg == "--force-color" {
+			forceColor = true
+		}
+	}
+
+	m.Ui = &cli.BasicUi{
+		Reader:      os.Stdin,
+		Writer:      colorable.NewColorableStdout(),
+		ErrorWriter: colorable.NewColorableStderr(),
+	}
+
+	// Only use colored UI if not disabled and stdout is a tty or colors are
+	// forced.
+	isTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
+	useColor := !noColor && (isTerminal || forceColor)
+	if useColor {
+		m.Ui = &cli.ColoredUi{
+			ErrorColor: cli.UiColorRed,
+			WarnColor:  cli.UiColorYellow,
+			InfoColor:  cli.UiColorGreen,
+			Ui:         m.Ui,
+		}
+	}
+}
+
 type usageOptsFlags uint8
 
 const (
@@ -196,12 +234,17 @@ func generalOptionsUsage(usageOpts usageOptsFlags) string {
 `
 
 	// note: that although very few commands use color explicitly, all of them
-	// return red-colored text on error so we don't want to make this
-	// configurable
+	// return red-colored text on error so we want the color flags to always be
+	// present in the help messages.
 	remainingText := `
   -no-color
     Disables colored command output. Alternatively, NOMAD_CLI_NO_COLOR may be
-    set.
+    set. This option takes precedence over -force-color.
+
+  -force-color
+    Forces colored command output. This can be used in cases where the usual
+    terminal detection fails. Alternatively, NOMAD_CLI_FORCE_COLOR may be set.
+    This option has no effect if -no-color is also used.
 
   -ca-cert=<path>
     Path to a PEM encoded CA cert file to use to verify the

command/meta_test.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"reflect"
 	"sort"
+	"strings"
 	"testing"
 
 	"github.com/kr/pty"
@@ -27,6 +28,7 @@ func TestMeta_FlagSet(t *testing.T) {
 			[]string{
 				"address",
 				"no-color",
+				"force-color",
 				"region",
 				"namespace",
 				"ca-cert",
@@ -81,11 +83,45 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "disable colors via CLI flag",
 			SetupFn: func(t *testing.T, m *Meta) {
-				m.Ui = &cli.ColoredUi{}
-
-				fs := m.FlagSet("colorize_test", FlagSetDefault)
-				err := fs.Parse([]string{"-no-color"})
-				assert.NoError(t, err)
+				m.SetupUi([]string{"-no-color"})
+			},
+			ExpectColor: false,
+		},
+		{
+			Name: "disable colors via env var",
+			SetupFn: func(t *testing.T, m *Meta) {
+				os.Setenv(EnvNomadCLINoColor, "1")
+				m.SetupUi([]string{})
+			},
+			ExpectColor: false,
+		},
+		{
+			Name: "force colors via CLI flag",
+			SetupFn: func(t *testing.T, m *Meta) {
+				m.SetupUi([]string{"-force-color"})
+			},
+			ExpectColor: true,
+		},
+		{
+			Name: "force colors via env var",
+			SetupFn: func(t *testing.T, m *Meta) {
+				os.Setenv(EnvNomadCLIForceColor, "1")
+				m.SetupUi([]string{})
+			},
+			ExpectColor: true,
+		},
+		{
+			Name: "no color take predecence over force color via CLI flag",
+			SetupFn: func(t *testing.T, m *Meta) {
+				m.SetupUi([]string{"-no-color", "-force-color"})
+			},
+			ExpectColor: false,
+		},
+		{
+			Name: "no color take predecence over force color via env var",
+			SetupFn: func(t *testing.T, m *Meta) {
+				os.Setenv(EnvNomadCLINoColor, "1")
+				m.SetupUi([]string{"-force-color"})
 			},
 			ExpectColor: false,
 		},
@@ -104,6 +140,14 @@ func TestMeta_Colorize(t *testing.T) {
 			defer func() { os.Stdout = oldStdout }()
 			os.Stdout = tty
 
+			// Make sure Nomad environment variables are clean.
+			for _, envVar := range os.Environ() {
+				if strings.HasPrefix(envVar, "NOMAD") {
+					k := strings.SplitN(envVar, "=", 2)[0]
+					os.Unsetenv(k)
+				}
+			}
+
 			// Run test case.
 			m := &Meta{}
 			if tc.SetupFn != nil {

main.go

@@ -19,10 +19,8 @@ import (
 
 	"github.com/hashicorp/nomad/command"
 	"github.com/hashicorp/nomad/version"
-	colorable "github.com/mattn/go-colorable"
 	"github.com/mitchellh/cli"
 	"github.com/sean-/seed"
-	"golang.org/x/crypto/ssh/terminal"
 )
 
 var (
@@ -88,24 +86,9 @@ func Run(args []string) int {
 }
 
 func RunCustom(args []string) int {
-	// Parse flags into env vars for global use
-	args = setupEnv(args)
-
 	// Create the meta object
 	metaPtr := new(command.Meta)
-
-	// Don't use color if disabled
-	color := true
-	if os.Getenv(command.EnvNomadCLINoColor) != "" {
-		color = false
-	}
-
-	isTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
-	metaPtr.Ui = &cli.BasicUi{
-		Reader:      os.Stdin,
-		Writer:      colorable.NewColorableStdout(),
-		ErrorWriter: colorable.NewColorableStderr(),
-	}
+	metaPtr.SetupUi(args)
 
 	// The Nomad agent never outputs color
 	agentUi := &cli.BasicUi{
@@ -114,16 +97,6 @@ func RunCustom(args []string) int {
 		ErrorWriter: os.Stderr,
 	}
 
-	// Only use colored UI if stdout is a tty, and not disabled
-	if isTerminal && color {
-		metaPtr.Ui = &cli.ColoredUi{
-			ErrorColor: cli.UiColorRed,
-			WarnColor:  cli.UiColorYellow,
-			InfoColor:  cli.UiColorGreen,
-			Ui:         metaPtr.Ui,
-		}
-	}
-
 	commands := command.Commands(metaPtr, agentUi)
 	cli := &cli.CLI{
 		Name:                       "nomad",
@@ -203,22 +176,3 @@ func printCommand(w io.Writer, name string, cmdFn cli.CommandFactory) {
 	}
 	fmt.Fprintf(w, "    %s\t%s\n", name, cmd.Synopsis())
 }
-
-// setupEnv parses args and may replace them and sets some env vars to known
-// values based on format options
-func setupEnv(args []string) []string {
-	noColor := false
-	for _, arg := range args {
-		// Check if color is set
-		if arg == "-no-color" || arg == "--no-color" {
-			noColor = true
-		}
-	}
-
-	// Put back into the env for later
-	if noColor {
-		os.Setenv(command.EnvNomadCLINoColor, "true")
-	}
-
-	return args
-}

website/content/partials/general_options.mdx

@@ -11,7 +11,12 @@
   user. Defaults to the "default" namespace.
 
 - `-no-color`: Disables colored command output. Alternatively,
-  `NOMAD_CLI_NO_COLOR` may be set.
+  `NOMAD_CLI_NO_COLOR` may be set. This option takes precedence over
+  `-force-color`.
+
+-`-force-color`: Forces colored command output. This can be used in cases where
+  the usual terminal detection fails. Alternatively, `NOMAD_CLI_FORCE_COLOR`
+  may be set. This option has no effect if `-no-color` is also used.
 
 - `-ca-cert=<path>`: Path to a PEM encoded CA cert file to use to verify the
   Nomad server SSL certificate. Overrides the `NOMAD_CACERT` environment

website/content/partials/general_options_no_namespace.mdx

@@ -6,7 +6,12 @@
   Agent's local region.
 
 - `-no-color`: Disables colored command output. Alternatively,
-  `NOMAD_CLI_NO_COLOR` may be set.
+  `NOMAD_CLI_NO_COLOR` may be set. This option takes precedence over
+  `-force-color`.
+
+-`-force-color`: Forces colored command output. This can be used in cases where
+  the usual terminal detection fails. Alternatively, `NOMAD_CLI_FORCE_COLOR`
+  may be set. This option has no effect if `-no-color` is also used.
 
 - `-ca-cert=<path>`: Path to a PEM encoded CA cert file to use to verify the
   Nomad server SSL certificate. Overrides the `NOMAD_CACERT` environment