fix: refactor janky code

coder · DanielleMaywood · Sep 17, 2024 · Sep 13, 2024 · Sep 13, 2024 · Sep 13, 2024
commit 40eb24ee2b9888d3f6c28ce14491117f6b6eaf81
diff --git a/cli/cliui/cliui.go b/cli/cliui/cliui.go
@@ -2,10 +2,7 @@ package cliui
 
 import (
 	"flag"
-	"fmt"
 	"os"
-	"slices"
-	"sync"
 	"time"
 
 	"github.com/muesli/termenv"
@@ -36,46 +33,9 @@ type Styles struct {
 }
 
 var (
-	color     termenv.Profile
-	colorOnce sync.Once
+	color termenv.Profile
 )
 
-var (
-	// ANSI color codes
-	red           = Color("1")
-	green         = Color("2")
-	yellow        = Color("3")
-	magenta       = Color("5")
-	white         = Color("7")
-	brightBlue    = Color("12")
-	brightMagenta = Color("13")
-)
-
-// Color returns a color for the given string.
-func Color(s string) termenv.Color {
-	colorOnce.Do(func() {
-		color = termenv.NewOutput(os.Stdout).ColorProfile()
-
-		if flag.Lookup("test.v") != nil {
-			// Use a consistent colorless profile in tests so that results
-			// are deterministic.
-			color = termenv.Ascii
-		}
-
-		// Currently it appears there is no way to use the flag from
-		// serpent as it isn't possible to create a root middleware that
-		// runs for every command. For now we just check if `os.Args`
-		// has the flag.
-		if slices.Contains(os.Args, fmt.Sprintf("--%s", NoColorFlag)) ||
-			slices.Contains(os.Args, fmt.Sprintf("--%s=true", NoColorFlag)) ||
-			os.Getenv("CODER_NO_COLOR") != "" ||
-			os.Getenv("NO_COLOR") != "" {
-			color = termenv.Ascii
-		}
-	})
-	return color.Color(s)
-}
-
 func isTerm() bool {
 	return color != termenv.Ascii
 }
@@ -135,7 +95,31 @@ func ifTerm(f pretty.Formatter) pretty.Formatter {
 	return f
 }
 
-func init() {
+type InitOptions struct {
+	NoColor bool
+}
+
+func Init(opts InitOptions) {
+	color = termenv.NewOutput(os.Stdout).ColorProfile()
+
+	if flag.Lookup("test.v") != nil {
+		// Use a consistent colorless profile in tests so that results
+		// are deterministic.
+		color = termenv.Ascii
+	}
+
+	if opts.NoColor {
+		color = termenv.Ascii
+	}
+
+	red := color.Color("1")
+	green := color.Color("2")
+	yellow := color.Color("3")
+	magenta := color.Color("5")
+	white := color.Color("7")
+	brightBlue := color.Color("12")
+	brightMagenta := color.Color("13")
+
 	// We do not adapt the color based on whether the terminal is light or dark.
 	// Doing so would require a round-trip between the program and the terminal
 	// due to the OSC query and response.

diff --git a/cli/root.go b/cli/root.go
@@ -331,12 +331,21 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 		r.clientURL = new(url.URL)
 	}
 
-	// NOTE(Danielle): It appears there is no way to have a 'global' middleware in
-	// 				   serpent so we manually handle the ENV/flag lookup and setup
-	//				   the option in the below OptionSet so it is documented.
-	//                 We use (and discard) this local variable to get the correct
-	//                 behavior from the CLI when the option is passed.
-	var noColorDiscarded bool
+	// Add a wrapper to every command to ensure we've loaded our CLI theme prior
+	// to running any handler.
+	var noColor bool
+	cmd.Walk(func(c *serpent.Command) {
+		middleware := func(next serpent.HandlerFunc) serpent.HandlerFunc {
+			cliui.Init(cliui.InitOptions{NoColor: noColor})
+			return next
+		}
+
+		if c.Middleware != nil {
+			c.Middleware = serpent.Chain(c.Middleware, middleware)
+		} else {
+			c.Middleware = middleware
+		}
+	})
 
 	globalGroup := &serpent.Group{
 		Name:        "Global",
@@ -474,7 +483,7 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 			Default:     "false",
 			Description: "Disable use of color in CLI output.",
 			Group:       globalGroup,
-			Value:       serpent.BoolOf(&noColorDiscarded),
+			Value:       serpent.BoolOf(&noColor),
 		},
 		{
 			Flag: "version",

diff --git a/cmd/cliui/main.go b/cmd/cliui/main.go
@@ -38,16 +38,6 @@ func main() {
 		},
 	}
 
-	var noColorDiscarded bool
-
-	root.Options = []serpent.Option{
-		{
-			Default: "false",
-			Flag:    cliui.NoColorFlag,
-			Value:   serpent.BoolOf(&noColorDiscarded),
-		},
-	}
-
 	root.Children = append(root.Children, &serpent.Command{
 		Use: "colors",
 		Handler: func(inv *serpent.Invocation) error {
@@ -398,6 +388,28 @@ func main() {
 		},
 	})
 
+	var noColor bool
+	root.Options = []serpent.Option{
+		{
+			Default: "false",
+			Flag:    cliui.NoColorFlag,
+			Value:   serpent.BoolOf(&noColor),
+		},
+	}
+
+	root.Walk(func(cmd *serpent.Command) {
+		middleware := func(next serpent.HandlerFunc) serpent.HandlerFunc {
+			cliui.Init(cliui.InitOptions{NoColor: noColor})
+			return next
+		}
+
+		if cmd.Middleware != nil {
+			cmd.Middleware = serpent.Chain(cmd.Middleware, middleware)
+		} else {
+			cmd.Middleware = middleware
+		}
+	})
+
 	err := root.Invoke(os.Args[1:]...).WithOS().Run()
 	if err != nil {
 		_, _ = fmt.Println(err.Error())