Add FlagVarDef and FlagVarSliceDef

2 files changed, 41 insertions, 1 deletions

diff --git a/flag.go b/flag.go
index bed32eb..ce394b8 100644
--- a/flag.go
+++ b/flag.go
@@ -15,6 +15,12 @@ func FlagVar[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T]) *T
 	return v.Value
 }
 
+func FlagVarDef[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T], def T) *T {
+	val := def
+	FlagVarPtr(fs, name, usage, parse, &val)
+	return &val
+}
+
 func FlagVarPtr[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T], val *T) {
 	fs.Var(&flagValue[T]{Parse: parse, Value: val}, name, usage)
 }
@@ -25,6 +31,13 @@ func FlagVarSlice[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T
 	return v.Values
 }
 
+func FlagVarSliceDef[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T], def []T) *[]T {
+	vals := make([]T, len(def))
+	copy(vals, def)
+	FlagVarSlicePtr(fs, name, usage, parse, &vals)
+	return &vals
+}
+
 func FlagVarSlicePtr[T any](fs *flag.FlagSet, name, usage string, parse ParseFunc[T], vals *[]T) {
 	fs.Var(&flagValueSlice[T]{Parse: parse, Values: vals}, name, usage)
 }
@@ -66,6 +79,8 @@ func (f *flagValue[T]) String() string {
 type flagValueSlice[T any] struct {
 	Parse  ParseFunc[T]
 	Values *[]T
+
+	shouldAppend bool
 }
 
 var _ flag.Value = &flagValueSlice[any]{}
@@ -75,7 +90,12 @@ func (f *flagValueSlice[T]) Set(s string) error {
 	if err != nil {
 		return err
 	}
-	*f.Values = append(*f.Values, val)
+	if f.shouldAppend {
+		*f.Values = append(*f.Values, val)
+	} else {
+		*f.Values = []T{val}
+		f.shouldAppend = true
+	}
 	return nil
 }
 
diff --git a/flag_test.go b/flag_test.go
index 56711d5..a4da69b 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -20,6 +20,16 @@ func TestFlagVar(s *testing.T) {
 	t.AssertEqual(true, *fl)
 }
 
+func TestFlagVarDef(s *testing.T) {
+	t := core.T{T: s}
+
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	fl := core.FlagVarDef(fs, "test", "", strconv.Atoi, 42)
+	t.AssertEqual(42, *fl)
+	t.AssertErrorIs(nil, fs.Parse([]string{"-test=1"}))
+	t.AssertEqual(1, *fl)
+}
+
 func TestFlagVarPtr(s *testing.T) {
 	t := core.T{T: s}
 
@@ -39,6 +49,16 @@ func TestFlagVarSlice(s *testing.T) {
 	t.AssertEqual([]int{1, 2, 42}, *fl)
 }
 
+func TestFlagVarSliceDef(s *testing.T) {
+	t := core.T{T: s}
+
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	fl := core.FlagVarSliceDef(fs, "test", "", strconv.Atoi, []int{42})
+	t.AssertEqual([]int{42}, *fl)
+	t.AssertErrorIs(nil, fs.Parse([]string{"-test=1", "-test=2"}))
+	t.AssertEqual([]int{1, 2}, *fl)
+}
+
 func TestFlagVarSlicePtr(s *testing.T) {
 	t := core.T{T: s}