Merge #150521

150521: sniffarg: simplify and add bool parsing r=tbg a=tbg

See commits. This is in preparation for adding plotted charts (as images) to the asim datadriven tests, for which I'll want to detect the `-rewrite` flag from outside the `datadriven` package.

Epic: CRDB-25222

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

pkg/testutils/sniffarg/BUILD.bazel

@@ -5,7 +5,10 @@ go_library(
     srcs = ["sniffarg.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/testutils/sniffarg",
     visibility = ["//visibility:public"],
-    deps = ["@com_github_spf13_pflag//:pflag"],
+    deps = [
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_spf13_pflag//:pflag",
+    ],
 )
 
 go_test(

pkg/testutils/sniffarg/sniffarg.go

@@ -9,44 +9,39 @@ import (
 	"os"
 	"regexp"
 
+	"github.com/cockroachdb/errors"
 	"github.com/spf13/pflag"
 )
 
 // DoEnv calls Do with os.Args[1:] as the first argument.
-func DoEnv(name string, out *string) error {
+func DoEnv(name string, out interface{}) error {
 	return Do(os.Args[1:], name, out)
 }
 
-// Do looks for the flags specified in `names` (no leading dashes) and
-// sets the corresponding `out` values to the values found in `args`. This only
-// works for string flags and in particular does not reliably work for
-// "presence" flags, such as bools, since these flags don't carry an explicit
-// value in the args.
+// Do looks for the flag `name` (no leading dashes) and sets the corresponding
+// `out` values to the value found in `args`.
+// Currently, `out` must be of type `*string` or `*bool`, though additional
+// types should be straightforward to add as needed.
 //
-// This is a helper for benchmarks that want to react to flags from their
-// environment.
-func Do(inArgs []string, name string, out *string) error {
+// This is a helper for tests and benchmarks that want to react to flags from
+// their environment.
+func Do(args []string, name string, out interface{}) error {
 	pf := pflag.NewFlagSet("test", pflag.ContinueOnError)
-	pf.StringVar(out, name, "", "")
-	var args []string
-	var addNext bool
-	for _, arg := range inArgs {
-		if addNext {
-			addNext = false
-			args = append(args, arg)
-		}
+	switch t := out.(type) {
+	case *string:
+		pf.StringVar(t, name, "", "")
+	case *bool:
+		pf.BoolVar(t, name, false, "")
+	default:
+		return errors.Errorf("unsupported type %T", t)
+	}
+	pf.ParseErrorsWhitelist = pflag.ParseErrorsWhitelist{UnknownFlags: true}
+	args = append([]string(nil), args...)
+	for i, arg := range args {
 		re := regexp.MustCompile(`^(-{1,2})` + regexp.QuoteMeta(name) + `(=|$)`)
-		if matches := re.FindStringSubmatch(arg); len(matches) > 0 {
-			if len(matches[1]) == 1 {
-				// Transform `-foo` into `--foo` for pflag-style flag.
-				arg = "-" + arg
-			}
-			if len(matches[2]) == 0 {
-				// The matched flag is of form `--foo bar` (vs `--foo=bar`), so value
-				// is next arg.
-				addNext = true
-			}
-			args = append(args, arg)
+		if matches := re.FindStringSubmatch(arg); len(matches) > 0 && len(matches[1]) == 1 {
+			// Transform `-foo` into `--foo` for pflag-style flag.
+			args[i] = "-" + arg
 		}
 	}
 	return pf.Parse(args)

pkg/testutils/sniffarg/sniffarg_test.go

@@ -21,18 +21,28 @@ func TestDo(t *testing.T) {
 		"-test.outputdir", "banana",
 		"something",
 		"--somethingelse", "foo",
-		"--boolflag",
+		"--falseflag=false",
+		"--trueflag",
 	}
 	var benchMem string
 	var outputDir string
 	var somethingElse string
+	var trueFlag bool
+	var falseFlag bool
+	var wrongType struct{}
 	notFound := "hello"
+	require.Error(t, Do(args, "test.benchmem", &wrongType))
 	require.NoError(t, Do(args, "test.benchmem", &benchMem))
 	require.NoError(t, Do(args, "test.outputdir", &outputDir))
 	require.NoError(t, Do(args, "somethingelse", &somethingElse))
 	require.NoError(t, Do(args, "notfound", &notFound))
+	require.NoError(t, Do(args, "falseflag", &falseFlag))
+	require.NoError(t, Do(args, "trueflag", &trueFlag))
 	assert.Equal(t, "5", benchMem)
 	assert.Equal(t, "banana", outputDir)
 	assert.Equal(t, "foo", somethingElse)
+	assert.False(t, falseFlag)
+	assert.True(t, trueFlag)
 	assert.Zero(t, notFound)
+
 }