Merge pull request ogier#20 from eparis/optional-args

Fix '--flag arg' and Allow flags to take optional arguments

bool.go

@@ -50,31 +50,31 @@ func (f *FlagSet) GetBool(name string) (bool, error) {
 // BoolVar defines a bool flag with specified name, default value, and usage string.
 // The argument p points to a bool variable in which to store the value of the flag.
 func (f *FlagSet) BoolVar(p *bool, name string, value bool, usage string) {
-	f.VarP(newBoolValue(value, p), name, "", usage)
+	f.BoolVarP(p, name, "", value, usage)
 }
 
 // Like BoolVar, but accepts a shorthand letter that can be used after a single dash.
 func (f *FlagSet) BoolVarP(p *bool, name, shorthand string, value bool, usage string) {
-	f.VarP(newBoolValue(value, p), name, shorthand, usage)
+	flag := f.VarPF(newBoolValue(value, p), name, shorthand, usage)
+	flag.NoOptDefVal = "true"
 }
 
 // BoolVar defines a bool flag with specified name, default value, and usage string.
 // The argument p points to a bool variable in which to store the value of the flag.
 func BoolVar(p *bool, name string, value bool, usage string) {
-	CommandLine.VarP(newBoolValue(value, p), name, "", usage)
+	BoolVarP(p, name, "", value, usage)
 }
 
 // Like BoolVar, but accepts a shorthand letter that can be used after a single dash.
 func BoolVarP(p *bool, name, shorthand string, value bool, usage string) {
-	CommandLine.VarP(newBoolValue(value, p), name, shorthand, usage)
+	flag := CommandLine.VarPF(newBoolValue(value, p), name, shorthand, usage)
+	flag.NoOptDefVal = "true"
 }
 
 // Bool defines a bool flag with specified name, default value, and usage string.
 // The return value is the address of a bool variable that stores the value of the flag.
 func (f *FlagSet) Bool(name string, value bool, usage string) *bool {
-	p := new(bool)
-	f.BoolVarP(p, name, "", value, usage)
-	return p
+	return f.BoolP(name, "", value, usage)
 }
 
 // Like Bool, but accepts a shorthand letter that can be used after a single dash.
@@ -87,10 +87,11 @@ func (f *FlagSet) BoolP(name, shorthand string, value bool, usage string) *bool
 // Bool defines a bool flag with specified name, default value, and usage string.
 // The return value is the address of a bool variable that stores the value of the flag.
 func Bool(name string, value bool, usage string) *bool {
-	return CommandLine.BoolP(name, "", value, usage)
+	return BoolP(name, "", value, usage)
 }
 
 // Like Bool, but accepts a shorthand letter that can be used after a single dash.
 func BoolP(name, shorthand string, value bool, usage string) *bool {
-	return CommandLine.BoolP(name, shorthand, value, usage)
+	b := CommandLine.BoolP(name, shorthand, value, usage)
+	return b
 }

bool_test.go

@@ -59,7 +59,8 @@ func (v *triStateValue) Type() string {
 func setUpFlagSet(tristate *triStateValue) *FlagSet {
 	f := NewFlagSet("test", ContinueOnError)
 	*tristate = triStateFalse
-	f.VarP(tristate, "tristate", "t", "tristate value (true, maybe or false)")
+	flag := f.VarPF(tristate, "tristate", "t", "tristate value (true, maybe or false)")
+	flag.NoOptDefVal = "true"
 	return f
 }
 
@@ -162,3 +163,18 @@ func TestInvalidValue(t *testing.T) {
 		t.Fatal("expected an error but did not get any, tristate has value", tristate)
 	}
 }
+
+func TestBoolP(t *testing.T) {
+	b := BoolP("bool", "b", false, "bool value in CommandLine")
+	c := BoolP("c", "c", false, "other bool value")
+	args := []string{"--bool"}
+	if err := CommandLine.Parse(args); err != nil {
+		t.Error("expected no error, got ", err)
+	}
+	if *b != true {
+		t.Errorf("expected b=true got b=%s", b)
+	}
+	if *c != false {
+		t.Errorf("expect c=false got c=%s", c)
+	}
+}

flag.go

@@ -152,6 +152,7 @@ type Flag struct {
 	Value       Value               // value as set
 	DefValue    string              // default value (as text); for usage message
 	Changed     bool                // If the user set the value (or if left to default)
+	NoOptDefVal string              //default value (as text); if the flag is on the command line without any options
 	Deprecated  string              // If this flag is deprecated, this string is the new or now thing to use
 	Annotations map[string][]string // used by cobra.Command  bash autocomple code
 }
@@ -341,16 +342,25 @@ func (f *FlagSet) PrintDefaults() {
 		if len(flag.Deprecated) > 0 {
 			return
 		}
-		format := "--%s=%s: %s\n"
-		if _, ok := flag.Value.(*stringValue); ok {
-			// put quotes on the value
-			format = "--%s=%q: %s\n"
-		}
+		format := ""
+		// ex: w/ option string argument '-%s, --%s[=%q]: %s\n'
 		if len(flag.Shorthand) > 0 {
-			format = "  -%s, " + format
+			format = "  -%s, --%s"
 		} else {
-			format = "   %s   " + format
+			format = "   %s   --%s"
+		}
+		if len(flag.NoOptDefVal) > 0 {
+			format = format + "["
+		}
+		if _, ok := flag.Value.(*stringValue); ok {
+			format = format + "=%q"
+		} else {
+			format = format + "=%s"
+		}
+		if len(flag.NoOptDefVal) > 0 {
+			format = format + "]"
 		}
+		format = format + ": %s\n"
 		fmt.Fprintf(f.out(), format, flag.Shorthand, flag.Name, flag.DefValue, flag.Usage)
 	})
 }
@@ -443,8 +453,8 @@ func (f *FlagSet) Var(value Value, name string, usage string) {
 	f.VarP(value, name, "", usage)
 }
 
-// Like Var, but accepts a shorthand letter that can be used after a single dash.
-func (f *FlagSet) VarP(value Value, name, shorthand, usage string) {
+// Like VarP, but returns the flag created
+func (f *FlagSet) VarPF(value Value, name, shorthand, usage string) *Flag {
 	// Remember the default value as a string; it won't change.
 	flag := &Flag{
 		Name:      name,
@@ -454,6 +464,12 @@ func (f *FlagSet) VarP(value Value, name, shorthand, usage string) {
 		DefValue:  value.String(),
 	}
 	f.AddFlag(flag)
+	return flag
+}
+
+// Like Var, but accepts a shorthand letter that can be used after a single dash.
+func (f *FlagSet) VarP(value Value, name, shorthand, usage string) {
+	_ = f.VarPF(value, name, shorthand, usage)
 }
 
 func (f *FlagSet) AddFlag(flag *Flag) {
@@ -566,11 +582,15 @@ func (f *FlagSet) parseLongArg(s string, args []string) (a []string, err error)
 	if len(split) == 2 {
 		// '--flag=arg'
 		value = split[1]
-	} else if bv, ok := flag.Value.(boolFlag); ok && bv.IsBoolFlag() {
-		// '--flag' (where flag is a bool)
-		value = "true"
+	} else if len(flag.NoOptDefVal) > 0 {
+		// '--flag' (arg was optional)
+		value = flag.NoOptDefVal
+	} else if len(a) > 0 {
+		// '--flag arg'
+		value = a[0]
+		a = a[1:]
 	} else {
-		// '--flag' (where flag was not a bool)
+		// '--flag' (arg was required)
 		err = f.failf("flag needs an argument: %s", s)
 		return
 	}
@@ -598,8 +618,8 @@ func (f *FlagSet) parseSingleShortArg(shorthands string, args []string) (outShor
 	if len(shorthands) > 2 && shorthands[1] == '=' {
 		value = shorthands[2:]
 		outShorts = ""
-	} else if bv, ok := flag.Value.(boolFlag); ok && bv.IsBoolFlag() {
-		value = "true"
+	} else if len(flag.NoOptDefVal) > 0 {
+		value = flag.NoOptDefVal
 	} else if len(shorthands) > 1 {
 		value = shorthands[1:]
 		outShorts = ""

flag_test.go

@@ -27,6 +27,7 @@ var (
 	test_string                  = String("test_string", "0", "string value")
 	test_float64                 = Float64("test_float64", 0, "float64 value")
 	test_duration                = Duration("test_duration", 0, "time.Duration value")
+	test_optional_int            = Int("test_optional_int", 0, "optional int value")
 	normalizeFlagNameInvocations = 0
 )
 
@@ -58,7 +59,7 @@ func TestEverything(t *testing.T) {
 		}
 	}
 	VisitAll(visitor)
-	if len(m) != 8 {
+	if len(m) != 9 {
 		t.Error("VisitAll misses some flags")
 		for k, v := range m {
 			t.Log(k, *v)
@@ -81,9 +82,10 @@ func TestEverything(t *testing.T) {
 	Set("test_string", "1")
 	Set("test_float64", "1")
 	Set("test_duration", "1s")
+	Set("test_optional_int", "1")
 	desired = "1"
 	Visit(visitor)
-	if len(m) != 8 {
+	if len(m) != 9 {
 		t.Error("Visit fails after set")
 		for k, v := range m {
 			t.Log(k, *v)
@@ -161,6 +163,10 @@ func testParse(f *FlagSet, t *testing.T) {
 	ipFlag := f.IP("ip", net.ParseIP("127.0.0.1"), "ip value")
 	maskFlag := f.IPMask("mask", ParseIPv4Mask("0.0.0.0"), "mask value")
 	durationFlag := f.Duration("duration", 5*time.Second, "time.Duration value")
+	optionalIntNoValueFlag := f.Int("optional-int-no-value", 0, "int value")
+	f.Lookup("optional-int-no-value").NoOptDefVal = "9"
+	optionalIntWithValueFlag := f.Int("optional-int-with-value", 0, "int value")
+	f.Lookup("optional-int-no-value").NoOptDefVal = "9"
 	extra := "one-extra-argument"
 	args := []string{
 		"--bool",
@@ -170,7 +176,7 @@ func testParse(f *FlagSet, t *testing.T) {
 		"--int8=-8",
 		"--int32=-32",
 		"--int64=0x23",
-		"--uint=24",
+		"--uint", "24",
 		"--uint8=8",
 		"--uint16=16",
 		"--uint32=32",
@@ -181,6 +187,8 @@ func testParse(f *FlagSet, t *testing.T) {
 		"--ip=10.11.12.13",
 		"--mask=255.255.255.0",
 		"--duration=2m",
+		"--optional-int-no-value",
+		"--optional-int-with-value=42",
 		extra,
 	}
 	if err := f.Parse(args); err != nil {
@@ -294,6 +302,12 @@ func testParse(f *FlagSet, t *testing.T) {
 	if _, err := f.GetInt("duration"); err == nil {
 		t.Error("GetInt parsed a time.Duration?!?!")
 	}
+	if *optionalIntNoValueFlag != 9 {
+		t.Error("optional int flag should be the default value, is ", *optionalIntNoValueFlag)
+	}
+	if *optionalIntWithValueFlag != 42 {
+		t.Error("optional int flag should be 42, is ", *optionalIntWithValueFlag)
+	}
 	if len(f.Args()) != 1 {
 		t.Error("expected one argument, got", len(f.Args()))
 	} else if f.Args()[0] != extra {