Merge pull request #19158 from vespa-engine/mpolden/color-control

Add flag/config option to control colorization

4 files changed, 56 insertions, 14 deletions

diff --git a/client/go/cmd/command_tester.go b/client/go/cmd/command_tester.go
index 095a1af7ac3..3d19e772875 100644
--- a/client/go/cmd/command_tester.go
+++ b/client/go/cmd/command_tester.go
@@ -8,7 +8,6 @@ import (
 	"bytes"
 	"crypto/tls"
 	"io/ioutil"
-	"log"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -16,7 +15,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/logrusorgru/aurora"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
@@ -34,9 +32,6 @@ func execute(cmd command, t *testing.T, client *mockHttpClient) string {
 		util.ActiveHttpClient = client
 	}
 
-	// Never print colors in tests
-	color = aurora.NewAurora(false)
-
 	// Set config dir. Use a separate one per test if none is specified
 	if cmd.homeDir == "" {
 		cmd.homeDir = t.TempDir()
@@ -62,7 +57,9 @@ func execute(cmd command, t *testing.T, client *mockHttpClient) string {
 
 	// Capture stdout and execute command
 	var b bytes.Buffer
-	log.SetOutput(&b)
+	stdout = &b
+
+	// Execute command and return output
 	rootCmd.SetArgs(append(cmd.args, cmd.moreArgs...))
 	rootCmd.Execute()
 	out, err := ioutil.ReadAll(&b)
diff --git a/client/go/cmd/config.go b/client/go/cmd/config.go
index 13142c92553..3753d9a9390 100644
--- a/client/go/cmd/config.go
+++ b/client/go/cmd/config.go
@@ -10,6 +10,7 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+	"sort"
 	"strconv"
 	"strings"
 
@@ -79,8 +80,14 @@ var getConfigCmd = &cobra.Command{
 		}
 
 		if len(args) == 0 { // Print all values
-			printOption(cfg, targetFlag)
-			printOption(cfg, applicationFlag)
+			var flags []string
+			for flag := range flagToConfigBindings {
+				flags = append(flags, flag)
+			}
+			sort.Strings(flags)
+			for _, flag := range flags {
+				printOption(cfg, flag)
+			}
 		} else {
 			printOption(cfg, args[0])
 		}
@@ -218,6 +225,12 @@ func (c *Config) Set(option, value string) error {
 		}
 		viper.Set(option, value)
 		return nil
+	case colorFlag:
+		switch value {
+		case "auto", "never", "always":
+			viper.Set(option, value)
+			return nil
+		}
 	}
 	return fmt.Errorf("invalid option or value: %q: %q", option, value)
 }
diff --git a/client/go/cmd/config_test.go b/client/go/cmd/config_test.go
index 07d165d58e0..70ad5d558e1 100644
--- a/client/go/cmd/config_test.go
+++ b/client/go/cmd/config_test.go
@@ -22,7 +22,7 @@ func TestConfig(t *testing.T) {
 	assert.Equal(t, "", execute(command{homeDir: homeDir, args: []string{"config", "set", "application", "t1.a1.i1"}}, t, nil))
 	assert.Equal(t, "application = t1.a1.i1\n", execute(command{homeDir: homeDir, args: []string{"config", "get", "application"}}, t, nil))
 
-	assert.Equal(t, "target = https://127.0.0.1\napplication = t1.a1.i1\n", execute(command{homeDir: homeDir, args: []string{"config", "get"}}, t, nil))
+	assert.Equal(t, "application = t1.a1.i1\ncolor = auto\ntarget = https://127.0.0.1\nwait = 0\n", execute(command{homeDir: homeDir, args: []string{"config", "get"}}, t, nil))
 
 	assert.Equal(t, "", execute(command{homeDir: homeDir, args: []string{"config", "set", "wait", "60"}}, t, nil))
 	assert.Equal(t, "wait option must be an integer >= 0, got \"foo\"\n", execute(command{homeDir: homeDir, args: []string{"config", "set", "wait", "foo"}}, t, nil))
diff --git a/client/go/cmd/root.go b/client/go/cmd/root.go
index d218d3639b1..4202035af92 100644
--- a/client/go/cmd/root.go
+++ b/client/go/cmd/root.go
@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"fmt"
 	"log"
 	"os"
 
@@ -27,34 +28,65 @@ Prefer web service API's to this in production.
 
 Vespa documentation: https://docs.vespa.ai`,
 		DisableAutoGenTag: true,
+		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+			configureOutput()
+		},
 	}
 
-	color          aurora.Aurora
 	targetArg      string
 	applicationArg string
 	waitSecsArg    int
+	colorArg       string
+
+	color  = aurora.NewAurora(false)
+	stdout = colorable.NewColorableStdout()
 )
 
 const (
 	applicationFlag = "application"
 	targetFlag      = "target"
 	waitFlag        = "wait"
+	colorFlag       = "color"
 )
 
-func configureLogger() {
-	color = aurora.NewAurora(isatty.IsTerminal(os.Stdout.Fd()))
+func configureOutput() {
 	log.SetFlags(0) // No timestamps
-	log.SetOutput(colorable.NewColorableStdout())
+	log.SetOutput(stdout)
+
+	config, err := LoadConfig()
+	if err != nil {
+		fatalErr(err, "Could not load config")
+	}
+	colorValue, err := config.Get(colorFlag)
+	if err != nil {
+		fatalErr(err)
+	}
+
+	colorize := false
+	switch colorValue {
+	case "auto":
+		file, ok := stdout.(*os.File)
+		if ok {
+			colorize = isatty.IsTerminal(file.Fd())
+		}
+	case "always":
+		colorize = true
+	case "never":
+	default:
+		fatalErrHint(fmt.Errorf("Invalid value for %s option", colorFlag), "Must be \"auto\", \"never\" or \"always\"")
+	}
+	color = aurora.NewAurora(colorize)
 }
 
 func init() {
-	configureLogger()
 	rootCmd.PersistentFlags().StringVarP(&targetArg, targetFlag, "t", "local", "The name or URL of the recipient of this command")
 	rootCmd.PersistentFlags().StringVarP(&applicationArg, applicationFlag, "a", "", "The application to manage")
 	rootCmd.PersistentFlags().IntVarP(&waitSecsArg, waitFlag, "w", 0, "Number of seconds to wait for a service to become ready")
+	rootCmd.PersistentFlags().StringVarP(&colorArg, colorFlag, "c", "auto", "Whether to use colors in output. Can be \"auto\", \"never\" or \"always\"")
 	bindFlagToConfig(targetFlag, rootCmd)
 	bindFlagToConfig(applicationFlag, rootCmd)
 	bindFlagToConfig(waitFlag, rootCmd)
+	bindFlagToConfig(colorFlag, rootCmd)
 }
 
 // Execute executes the root command.