fix: allow overridding default string array

coder · ammario · Mar 30, 2023 · Mar 29, 2023 · Mar 29, 2023 · Mar 29, 2023
commit 5a03a5d17b1813ab47e78b59fec6d8d8563ddd05
diff --git a/cli/clibase/cmd.go b/cli/clibase/cmd.go
@@ -219,29 +219,7 @@ func copyFlagSetWithout(fs *pflag.FlagSet, without string) *pflag.FlagSet {
 // allArgs is wired through the stack so that global flags can be accepted
 // anywhere in the command invocation.
 func (i *Invocation) run(state *runState) error {
-	err := i.Command.Options.SetDefaults()
-	if err != nil {
-		return xerrors.Errorf("setting defaults: %w", err)
-	}
-
-	// If we set the Default of an array but later see a flag for it, we
-	// don't want to append, we want to replace. So, we need to keep the state
-	// of defaulted array options.
-	defaultedArrays := make(map[string]int)
-	for _, opt := range i.Command.Options {
-		sv, ok := opt.Value.(pflag.SliceValue)
-		if !ok {
-			continue
-		}
-
-		if opt.Flag == "" {
-			continue
-		}
-
-		defaultedArrays[opt.Flag] = len(sv.GetSlice())
-	}
-
-	err = i.Command.Options.ParseEnv(i.Environ)
+	err := i.Command.Options.ParseEnv(i.Environ)
 	if err != nil {
 		return xerrors.Errorf("parsing env: %w", err)
 	}
@@ -282,34 +260,24 @@ func (i *Invocation) run(state *runState) error {
 		// so we check the error after looking for a child command.
 		state.flagParseErr = i.parsedFlags.Parse(state.allArgs)
 		parsedArgs = i.parsedFlags.Args()
+	}
 
-		i.parsedFlags.VisitAll(func(f *pflag.Flag) {
-			i, ok := defaultedArrays[f.Name]
-			if !ok {
-				return
-			}
-
-			if !f.Changed {
-				return
-			}
-
-			// If flag was changed, we need to remove the default values.
-			sv, ok := f.Value.(pflag.SliceValue)
-			if !ok {
-				panic("defaulted array option is not a slice value")
-			}
-			ss := sv.GetSlice()
-			if len(ss) == 0 {
-				// Slice likely zeroed by a flag.
-				// E.g. "--fruit" may default to "apples,oranges" but the user
-				// provided "--fruit=""".
-				return
-			}
-			err := sv.Replace(ss[i:])
-			if err != nil {
-				panic(err)
-			}
-		})
+	// Set defaults for flags that weren't set by the user.
+	skipDefaults := make(map[string]struct{}, len(i.Command.Options))
+	i.parsedFlags.VisitAll(func(f *pflag.Flag) {
+		if !f.Changed {
+			return
+		}
+		skipDefaults[f.Name] = struct{}{}
+	})
+	for _, opt := range i.Command.Options {
+		if opt.envSet {
+			skipDefaults[opt.Name] = struct{}{}
+		}
+	}
+	err = i.Command.Options.SetDefaults(skipDefaults)
+	if err != nil {
+		return xerrors.Errorf("setting defaults: %w", err)
 	}
 
 	// Run child command if found (next child only)

diff --git a/cli/clibase/cmd_test.go b/cli/clibase/cmd_test.go
@@ -247,6 +247,7 @@ func TestCommand_FlagOverride(t *testing.T) {
 		Use: "1",
 		Options: clibase.OptionSet{
 			{
+				Name:  "flag",
 				Flag:  "f",
 				Value: clibase.DiscardValue,
 			},
@@ -256,6 +257,7 @@ func TestCommand_FlagOverride(t *testing.T) {
 				Use: "2",
 				Options: clibase.OptionSet{
 					{
+						Name:  "flag",
 						Flag:  "f",
 						Value: clibase.StringOf(&flag),
 					},
@@ -527,11 +529,17 @@ func TestCommand_EmptySlice(t *testing.T) {
 		}
 	}
 
-	// Base-case
+	// Base-case, uses default.
 	err := cmd("bad", "bad", "bad").Invoke().Run()
 	require.NoError(t, err)
 
+	// Reset to nothing at all.
 	inv := cmd().Invoke("--arr", "")
 	err = inv.Run()
 	require.NoError(t, err)
+
+	// Override
+	inv = cmd("great").Invoke("--arr", "great")
+	err = inv.Run()
+	require.NoError(t, err)
 }
diff --git a/cli/clibase/option.go b/cli/clibase/option.go
@@ -46,6 +46,8 @@ type Option struct {
 	UseInstead []Option `json:"use_instead,omitempty"`
 
 	Hidden bool `json:"hidden,omitempty"`
+
+	envSet bool
 }
 
 // OptionSet is a group of options that can be applied to a command.
@@ -133,6 +135,7 @@ func (s *OptionSet) ParseEnv(vs []EnvVar) error {
 			continue
 		}
 
+		opt.envSet = true
 		if err := opt.Value.Set(envVal); err != nil {
 			merr = multierror.Append(
 				merr, xerrors.Errorf("parse %q: %w", opt.Name, err),
@@ -145,17 +148,28 @@ func (s *OptionSet) ParseEnv(vs []EnvVar) error {
 
 // SetDefaults sets the default values for each Option.
 // It should be called before all parsing (e.g. ParseFlags, ParseEnv).
-func (s *OptionSet) SetDefaults() error {
+func (s *OptionSet) SetDefaults(skip map[string]struct{}) error {
 	if s == nil {
 		return nil
 	}
 
 	var merr *multierror.Error
 
 	for _, opt := range *s {
+		if opt.Name == "" {
+			merr = multierror.Append(
+				merr, xerrors.Errorf("parse: no Name field set"),
+			)
+			continue
+		}
+		if _, ok := skip[opt.Name]; ok {
+			continue
+		}
+
 		if opt.Default == "" {
 			continue
 		}
+
 		if opt.Value == nil {
 			merr = multierror.Append(
 				merr,

diff --git a/cli/clibase/option_test.go b/cli/clibase/option_test.go
@@ -49,7 +49,7 @@ func TestOptionSet_ParseFlags(t *testing.T) {
 			},
 		}
 
-		err := os.SetDefaults()
+		err := os.SetDefaults(nil)
 		require.NoError(t, err)
 
 		err = os.FlagSet().Parse([]string{"--name", "foo", "--name", "bar"})
@@ -111,7 +111,7 @@ func TestOptionSet_ParseEnv(t *testing.T) {
 			},
 		}
 
-		err := os.SetDefaults()
+		err := os.SetDefaults(nil)
 		require.NoError(t, err)
 
 		err = os.ParseEnv(clibase.ParseEnviron([]string{"CODER_WORKSPACE_NAME="}, "CODER_"))

diff --git a/cli/clibase/yaml_test.go b/cli/clibase/yaml_test.go
@@ -44,7 +44,7 @@ func TestOption_ToYAML(t *testing.T) {
 			},
 		}
 
-		err := os.SetDefaults()
+		err := os.SetDefaults(nil)
 		require.NoError(t, err)
 
 		n, err := os.ToYAML()

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -1075,7 +1075,7 @@ QastnN77KfUwdj3SJt44U/uh1jAIv4oSLBr8HYUkbnI8
 func DeploymentValues(t *testing.T) *codersdk.DeploymentValues {
 	var cfg codersdk.DeploymentValues
 	opts := cfg.Options()
-	err := opts.SetDefaults()
+	err := opts.SetDefaults(nil)
 	require.NoError(t, err)
 	return &cfg
 }