Add colorized code output flag in CLI (#752)

NethermindEth · May 16, 2023 · 66fdc74 · 66fdc74
1 parent 63b7bd9
commit 66fdc74
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 7 deletions.
diff --git a/cmd/juno/juno.go b/cmd/juno/juno.go
@@ -31,11 +31,13 @@ const (
 	dbPathF   = "db-path"
 	networkF  = "network"
 	pprofF    = "pprof"
+	colourF   = "colour"
 
 	defaultConfig  = ""
 	defaultRPCPort = uint16(6060)
 	defaultDBPath  = ""
 	defaultPprof   = false
+	defaultColour  = true
 
 	configFlagUsage   = "The yaml configuration file."
 	logLevelFlagUsage = "Options: debug, info, warn, error."
@@ -44,6 +46,7 @@ const (
 	dbPathUsage  = "Location of the database files."
 	networkUsage = "Options: mainnet, goerli, goerli2, integration."
 	pprofUsage   = "Enables the pprof server and listens on port 9080."
+	colourUsage  = "Uses --colour=false command to disable colourized outputs (ANSI Escape Codes)."
 )
 
 var Version string
@@ -76,7 +79,7 @@ func main() {
 	}
 }
 
-// NewCmd returns a command that can be exected with any of the Cobra Execute* functions.
+// NewCmd returns a command that can be executed with any of the Cobra Execute* functions.
 // The RunE field is set to the user-provided run function, allowing for robust testing setups.
 //
 //  1. NewCmd is called with a non-nil config and a run function.
@@ -125,6 +128,7 @@ func NewCmd(config *node.Config, run func(*cobra.Command, []string) error) *cobr
 	junoCmd.Flags().String(dbPathF, defaultDBPath, dbPathUsage)
 	junoCmd.Flags().Var(&defaultNetwork, networkF, networkUsage)
 	junoCmd.Flags().Bool(pprofF, defaultPprof, pprofUsage)
+	junoCmd.Flags().Bool(colourF, defaultColour, colourUsage)
 
 	return junoCmd
 }
diff --git a/cmd/juno/juno_test.go b/cmd/juno/juno_test.go
@@ -25,6 +25,7 @@ func TestConfigPrecedence(t *testing.T) {
 	defaultDBPath := ""
 	defaultNetwork := utils.MAINNET
 	defaultPprof := false
+	defaultColour := true
 
 	tests := map[string]struct {
 		cfgFile         bool
@@ -41,6 +42,7 @@ func TestConfigPrecedence(t *testing.T) {
 				DatabasePath: defaultDBPath,
 				Network:      defaultNetwork,
 				Pprof:        defaultPprof,
+				Colour:       defaultColour,
 			},
 		},
 		"config file path is empty string": {
@@ -51,6 +53,7 @@ func TestConfigPrecedence(t *testing.T) {
 				DatabasePath: defaultDBPath,
 				Network:      defaultNetwork,
 				Pprof:        defaultPprof,
+				Colour:       defaultColour,
 			},
 		},
 		"config file doesn't exist": {
@@ -64,6 +67,7 @@ func TestConfigPrecedence(t *testing.T) {
 				LogLevel: defaultLogLevel,
 				RPCPort:  defaultRPCPort,
 				Network:  defaultNetwork,
+				Colour:   defaultColour,
 			},
 		},
 		"config file with all settings but without any other flags": {
@@ -80,6 +84,7 @@ pprof: true
 				DatabasePath: "/home/.juno",
 				Network:      utils.GOERLI2,
 				Pprof:        true,
+				Colour:       defaultColour,
 			},
 		},
 		"config file with some settings but without any other flags": {
@@ -93,6 +98,7 @@ rpc-port: 4576
 				DatabasePath: defaultDBPath,
 				Network:      defaultNetwork,
 				Pprof:        defaultPprof,
+				Colour:       defaultColour,
 			},
 		},
 		"all flags without config file": {
@@ -106,6 +112,7 @@ rpc-port: 4576
 				DatabasePath: "/home/.juno",
 				Network:      utils.GOERLI,
 				Pprof:        true,
+				Colour:       defaultColour,
 			},
 		},
 		"some flags without config file": {
@@ -118,6 +125,7 @@ rpc-port: 4576
 				RPCPort:      4576,
 				DatabasePath: "/home/.juno",
 				Network:      utils.INTEGRATION,
+				Colour:       defaultColour,
 			},
 		},
 		"all setting set in both config file and flags": {
@@ -138,6 +146,7 @@ pprof: true
 				DatabasePath: "/home/flag/.juno",
 				Network:      utils.INTEGRATION,
 				Pprof:        true,
+				Colour:       defaultColour,
 			},
 		},
 		"some setting set in both config file and flags": {
@@ -153,6 +162,7 @@ network: goerli
 				DatabasePath: "/home/flag/.juno",
 				Network:      utils.GOERLI,
 				Pprof:        defaultPprof,
+				Colour:       defaultColour,
 			},
 		},
 		"some setting set in default, config file and flags": {
@@ -165,6 +175,7 @@ network: goerli
 				DatabasePath: "/home/flag/.juno",
 				Network:      utils.GOERLI2,
 				Pprof:        true,
+				Colour:       defaultColour,
 			},
 		},
 	}

diff --git a/node/node.go b/node/node.go
@@ -31,6 +31,7 @@ type Config struct {
 	DatabasePath string         `mapstructure:"db-path"`
 	Network      utils.Network  `mapstructure:"network"`
 	Pprof        bool           `mapstructure:"pprof"`
+	Colour       bool           `mapstructure:"colour"`
 }
 
 type Node struct {
@@ -52,7 +53,7 @@ func New(cfg *Config) (*Node, error) {
 		}
 		cfg.DatabasePath = filepath.Join(dirPrefix, cfg.Network.String())
 	}
-	log, err := utils.NewZapLogger(cfg.LogLevel)
+	log, err := utils.NewZapLogger(cfg.LogLevel, cfg.Colour)
 	if err != nil {
 		return nil, err
 	}
@@ -149,7 +150,7 @@ func makeHTTP(port uint16, rpcHandler *rpc.Handler, log utils.SimpleLogger) *jso
 func (n *Node) Run(ctx context.Context) {
 	n.log.Infow("Starting Juno...", "config", fmt.Sprintf("%+v", *n.cfg))
 
-	dbLog, err := utils.NewZapLogger(utils.ERROR)
+	dbLog, err := utils.NewZapLogger(utils.ERROR, n.cfg.Colour)
 	if err != nil {
 		n.log.Errorw("Error creating DB logger", "err", err)
 		return

diff --git a/utils/log.go b/utils/log.go
@@ -91,10 +91,13 @@ func NewNopZapLogger() *ZapLogger {
 	return &ZapLogger{zap.NewNop().Sugar()}
 }
 
-func NewZapLogger(logLevel LogLevel) (*ZapLogger, error) {
+func NewZapLogger(logLevel LogLevel, colour bool) (*ZapLogger, error) {
 	config := zap.NewProductionConfig()
 	config.Encoding = "console"
-	config.EncoderConfig.EncodeLevel = zapcore.CapitalLevelEncoder
+	config.EncoderConfig.EncodeLevel = zapcore.CapitalColorLevelEncoder
+	if !colour {
+		config.EncoderConfig.EncodeLevel = zapcore.CapitalLevelEncoder
+	}
 	config.EncoderConfig.EncodeTime = func(t time.Time, enc zapcore.PrimitiveArrayEncoder) {
 		enc.AppendString(t.Local().Format("15:04:05.000 02/01/2006 -07:00"))
 	}

diff --git a/utils/log_test.go b/utils/log_test.go
@@ -76,10 +76,19 @@ func TestLogLevelType(t *testing.T) {
 	assert.Equal(t, "LogLevel", new(utils.LogLevel).Type())
 }
 
-func TestZap(t *testing.T) {
+func TestZapWithColour(t *testing.T) {
 	for level, str := range levelStrings {
 		t.Run("level: "+str, func(t *testing.T) {
-			_, err := utils.NewZapLogger(level)
+			_, err := utils.NewZapLogger(level, true)
+			assert.NoError(t, err)
+		})
+	}
+}
+
+func TestZapWithoutColour(t *testing.T) {
+	for level, str := range levelStrings {
+		t.Run("level: "+str, func(t *testing.T) {
+			_, err := utils.NewZapLogger(level, false)
 			assert.NoError(t, err)
 		})
 	}