Make command line parsing available outside pflag

We basically want to allow callers to use our command line parser. There
is a user who wants to log the command line, but they don't want to log
"sensitive" flags. This allows that user to parse the command line and
get each flag. They can use information in the flag to write their own
printer.

2 files changed, 110 insertions, 10 deletions

diff --git a/flag.go b/flag.go
index fa81564..b4bfecb 100644
--- a/flag.go
+++ b/flag.go
@@ -752,7 +752,7 @@ func containsShorthand(arg, shorthand string) bool {
 	return strings.Contains(arg, shorthand)
 }
 
-func (f *FlagSet) parseLongArg(s string, args []string) (a []string, err error) {
+func (f *FlagSet) parseLongArg(s string, args []string, fn parseFunc) (a []string, err error) {
 	a = args
 	name := s[2:]
 	if len(name) == 0 || name[0] == '-' || name[0] == '=' {
@@ -786,11 +786,11 @@ func (f *FlagSet) parseLongArg(s string, args []string) (a []string, err error)
 		err = f.failf("flag needs an argument: %s", s)
 		return
 	}
-	err = f.setFlag(flag, value, s)
+	err = fn(flag, value, s)
 	return
 }
 
-func (f *FlagSet) parseSingleShortArg(shorthands string, args []string) (outShorts string, outArgs []string, err error) {
+func (f *FlagSet) parseSingleShortArg(shorthands string, args []string, fn parseFunc) (outShorts string, outArgs []string, err error) {
 	if strings.HasPrefix(shorthands, "test.") {
 		return
 	}
@@ -825,16 +825,16 @@ func (f *FlagSet) parseSingleShortArg(shorthands string, args []string) (outShor
 		err = f.failf("flag needs an argument: %q in -%s", c, shorthands)
 		return
 	}
-	err = f.setFlag(flag, value, shorthands)
+	err = fn(flag, value, shorthands)
 	return
 }
 
-func (f *FlagSet) parseShortArg(s string, args []string) (a []string, err error) {
+func (f *FlagSet) parseShortArg(s string, args []string, fn parseFunc) (a []string, err error) {
 	a = args
 	shorthands := s[1:]
 
 	for len(shorthands) > 0 {
-		shorthands, a, err = f.parseSingleShortArg(shorthands, args)
+		shorthands, a, err = f.parseSingleShortArg(shorthands, args, fn)
 		if err != nil {
 			return
 		}
@@ -843,7 +843,7 @@ func (f *FlagSet) parseShortArg(s string, args []string) (a []string, err error)
 	return
 }
 
-func (f *FlagSet) parseArgs(args []string) (err error) {
+func (f *FlagSet) parseArgs(args []string, fn parseFunc) (err error) {
 	for len(args) > 0 {
 		s := args[0]
 		args = args[1:]
@@ -863,9 +863,9 @@ func (f *FlagSet) parseArgs(args []string) (err error) {
 				f.args = append(f.args, args...)
 				break
 			}
-			args, err = f.parseLongArg(s, args)
+			args, err = f.parseLongArg(s, args, fn)
 		} else {
-			args, err = f.parseShortArg(s, args)
+			args, err = f.parseShortArg(s, args, fn)
 		}
 		if err != nil {
 			return
@@ -881,7 +881,41 @@ func (f *FlagSet) parseArgs(args []string) (err error) {
 func (f *FlagSet) Parse(arguments []string) error {
 	f.parsed = true
 	f.args = make([]string, 0, len(arguments))
-	err := f.parseArgs(arguments)
+
+	assign := func(flag *Flag, value, origArg string) error {
+		return f.setFlag(flag, value, origArg)
+	}
+
+	err := f.parseArgs(arguments, assign)
+	if err != nil {
+		switch f.errorHandling {
+		case ContinueOnError:
+			return err
+		case ExitOnError:
+			os.Exit(2)
+		case PanicOnError:
+			panic(err)
+		}
+	}
+	return nil
+}
+
+type parseFunc func(flag *Flag, value, origArg string) error
+
+// ParseAll parses flag definitions from the argument list, which should not
+// include the command name. The arguments for fn are flag and value. Must be
+// called after all flags in the FlagSet are defined and before flags are
+// accessed by the program. The return value will be ErrHelp if -help was set
+// but not defined.
+func (f *FlagSet) ParseAll(arguments []string, fn func(flag *Flag, value string) error) error {
+	f.parsed = true
+	f.args = make([]string, 0, len(arguments))
+
+	assign := func(flag *Flag, value, origArg string) error {
+		return fn(flag, value)
+	}
+
+	err := f.parseArgs(arguments, assign)
 	if err != nil {
 		switch f.errorHandling {
 		case ContinueOnError:
@@ -907,6 +941,14 @@ func Parse() {
 	CommandLine.Parse(os.Args[1:])
 }
 
+// ParseAll parses the command-line flags from os.Args[1:] and called fn for each.
+// The arguments for fn are flag and value. Must be called after all flags are
+// defined and before flags are accessed by the program.
+func ParseAll(fn func(flag *Flag, value string) error) {
+	// Ignore errors; CommandLine is set for ExitOnError.
+	CommandLine.ParseAll(os.Args[1:], fn)
+}
+
 // SetInterspersed sets whether to support interspersed option/non-option arguments.
 func SetInterspersed(interspersed bool) {
 	CommandLine.SetInterspersed(interspersed)
diff --git a/flag_test.go b/flag_test.go
index b294fc7..ad050b3 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -333,6 +333,59 @@ func testParse(f *FlagSet, t *testing.T) {
 	}
 }
 
+func testParseAll(f *FlagSet, t *testing.T) {
+	if f.Parsed() {
+		fmt.Errorf("f.Parse() = true before Parse")
+	}
+	f.BoolP("boola", "a", false, "bool value")
+	f.BoolP("boolb", "b", false, "bool2 value")
+	f.BoolP("boolc", "c", false, "bool3 value")
+	f.BoolP("boold", "d", false, "bool4 value")
+	f.StringP("stringa", "s", "0", "string value")
+	f.StringP("stringz", "z", "0", "string value")
+	f.StringP("stringx", "x", "0", "string value")
+	f.StringP("stringy", "y", "0", "string value")
+	f.Lookup("stringx").NoOptDefVal = "1"
+	args := []string{
+		"-ab",
+		"-cs=xx",
+		"--stringz=something",
+		"-d=true",
+		"-x",
+		"-y",
+		"ee",
+	}
+	want := []string{
+		"boola", "true",
+		"boolb", "true",
+		"boolc", "true",
+		"stringa", "xx",
+		"stringz", "something",
+		"boold", "true",
+		"stringx", "1",
+		"stringy", "ee",
+	}
+	got := []string{}
+	store := func(flag *Flag, value string) error {
+		got = append(got, flag.Name)
+		if len(value) > 0 {
+			got = append(got, value)
+		}
+		return nil
+	}
+	if err := f.ParseAll(args, store); err != nil {
+		t.Errorf("expected no error, got ", err)
+	}
+	if !f.Parsed() {
+		t.Errorf("f.Parse() = false after Parse")
+	}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("f.ParseAll() fail to restore the args")
+		t.Errorf("Got: %v", got)
+		t.Errorf("Want: %v", want)
+	}
+}
+
 func TestShorthand(t *testing.T) {
 	f := NewFlagSet("shorthand", ContinueOnError)
 	if f.Parsed() {
@@ -398,6 +451,11 @@ func TestParse(t *testing.T) {
 	testParse(GetCommandLine(), t)
 }
 
+func TestParseAll(t *testing.T) {
+	ResetForTesting(func() { t.Error("bad parse") })
+	testParseAll(GetCommandLine(), t)
+}
+
 func TestFlagSetParse(t *testing.T) {
 	testParse(NewFlagSet("test", ContinueOnError), t)
 }