script: Use pflag for command arguments

Embed the pflag FlagSet into the command creation
so we can show decent help for the flags.

Signed-off-by: Jussi Maki <[email protected]>

script/cmds.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/cilium/hive/script/internal/diff"
+	"github.com/spf13/pflag"
 	"golang.org/x/term"
 )
 
@@ -176,14 +177,19 @@ func Chmod() Cmd {
 		})
 }
 
+func compareFlags(fs *pflag.FlagSet) {
+	fs.BoolP("quiet", "q", false, "Suppress printing of diff")
+}
+
 // Cmp compares the contents of two files, or the contents of either the
 // "stdout" or "stderr" buffer and a file, returning a non-nil error if the
 // contents differ.
 func Cmp() Cmd {
 	return Command(
 		CmdUsage{
-			Args:    "[-q] file1 file2",
+			Args:    "file1 file2",
 			Summary: "compare files for differences",
+			Flags:   compareFlags,
 			Detail: []string{
 				"By convention, file1 is the actual data and file2 is the expected data.",
 				"The command succeeds if the file contents are identical.",
@@ -200,7 +206,8 @@ func Cmp() Cmd {
 func Cmpenv() Cmd {
 	return Command(
 		CmdUsage{
-			Args:    "[-q] file1 file2",
+			Args:    "file1 file2",
+			Flags:   compareFlags,
 			Summary: "compare files for differences, with environment expansion",
 			Detail: []string{
 				"By convention, file1 is the actual data and file2 is the expected data.",
@@ -214,10 +221,9 @@ func Cmpenv() Cmd {
 }
 
 func doCompare(s *State, env bool, args ...string) error {
-	quiet := false
-	if len(args) > 0 && args[0] == "-q" {
-		quiet = true
-		args = args[1:]
+	quiet, err := s.Flags.GetBool("quiet")
+	if err != nil {
+		return err
 	}
 	if len(args) != 2 {
 		return ErrUsage
@@ -577,23 +583,22 @@ func Exists() Cmd {
 	return Command(
 		CmdUsage{
 			Summary: "check that files exist",
-			Args:    "[-readonly] [-exec] file...",
+			Args:    "file...",
+			Flags: func(fs *pflag.FlagSet) {
+				fs.Bool("readonly", false, "File must not be writable")
+				fs.Bool("exec", false, "File must not be executable")
+			},
 		},
 		func(s *State, args ...string) (WaitFunc, error) {
-			var readonly, exec bool
-		loop:
-			for len(args) > 0 {
-				switch args[0] {
-				case "-readonly":
-					readonly = true
-					args = args[1:]
-				case "-exec":
-					exec = true
-					args = args[1:]
-				default:
-					break loop
-				}
+			readonly, err := s.Flags.GetBool("readonly")
+			if err != nil {
+				return nil, err
 			}
+			exec, err := s.Flags.GetBool("exec")
+			if err != nil {
+				return nil, err
+			}
+
 			if len(args) == 0 {
 				return nil, ErrUsage
 			}
@@ -625,7 +630,8 @@ func Grep() Cmd {
 	return Command(
 		CmdUsage{
 			Summary: "find lines in a file that match a pattern",
-			Args:    matchUsage + " file",
+			Args:    "'pattern' file",
+			Flags:   matchFlags,
 			Detail: []string{
 				"The command succeeds if at least one match (or the exact count, if given) is found.",
 				"The -q flag suppresses printing of matches.",
@@ -637,26 +643,20 @@ func Grep() Cmd {
 		})
 }
 
-const matchUsage = "[-count=N] [-q] 'pattern'"
+func matchFlags(fs *pflag.FlagSet) {
+	fs.Int("count", 0, "Exact count of matches")
+	fs.BoolP("quiet", "q", false, "Suppress printing of matches")
+}
 
 // match implements the Grep, Stdout, and Stderr commands.
 func match(s *State, args []string, text, name string) error {
-	n := 0
-	if len(args) >= 1 && strings.HasPrefix(args[0], "-count=") {
-		var err error
-		n, err = strconv.Atoi(args[0][len("-count="):])
-		if err != nil {
-			return fmt.Errorf("bad -count=: %v", err)
-		}
-		if n < 1 {
-			return fmt.Errorf("bad -count=: must be at least 1")
-		}
-		args = args[1:]
+	n, err := s.Flags.GetInt("count")
+	if err != nil {
+		return err
 	}
-	quiet := false
-	if len(args) >= 1 && args[0] == "-q" {
-		quiet = true
-		args = args[1:]
+	quiet, err := s.Flags.GetBool("quiet")
+	if err != nil {
+		return err
 	}
 
 	isGrep := name == "grep"
@@ -1008,7 +1008,8 @@ func Stderr() Cmd {
 	return Command(
 		CmdUsage{
 			Summary: "find lines in the stderr buffer that match a pattern",
-			Args:    matchUsage + " file",
+			Args:    "'pattern'",
+			Flags:   matchFlags,
 			Detail: []string{
 				"The command succeeds if at least one match (or the exact count, if given) is found.",
 				"The -q flag suppresses printing of matches.",
@@ -1025,7 +1026,8 @@ func Stdout() Cmd {
 	return Command(
 		CmdUsage{
 			Summary: "find lines in the stdout buffer that match a pattern",
-			Args:    matchUsage + " file",
+			Args:    "'pattern'",
+			Flags:   matchFlags,
 			Detail: []string{
 				"The command succeeds if at least one match (or the exact count, if given) is found.",
 				"The -q flag suppresses printing of matches.",

script/engine.go

@@ -62,6 +62,8 @@ import (
 	"sort"
 	"strings"
 	"time"
+
+	"github.com/spf13/pflag"
 )
 
 // An Engine stores the configuration for executing a set of scripts.
@@ -116,9 +118,17 @@ type WaitFunc func(*State) (stdout, stderr string, err error)
 // A CmdUsage describes the usage of a Cmd, independent of its name
 // (which can change based on its registration).
 type CmdUsage struct {
-	Summary string   // in the style of the Name section of a Unix 'man' page, omitting the name
-	Args    string   // a brief synopsis of the command's arguments (only)
-	Detail  []string // zero or more sentences in the style of the Description section of a Unix 'man' page
+	Summary string // in the style of the Name section of a Unix 'man' page, omitting the name
+
+	// synopsis of arguments, e.g. "files...". If [Flags] is provided then these will be prepended
+	// to [Args] in usage output.
+	Args string
+
+	// flags for the command, optional. If provided, then args passed to the command will not include any
+	// of the flag arguments. The [plafg.FlagSet] is available to the command via [State.Flags].
+	Flags func(fs *pflag.FlagSet)
+
+	Detail []string // zero or more sentences in the style of the Description section of a Unix 'man' page
 
 	// If Async is true, the Cmd is meaningful to run in the background, and its
 	// Run method must return either a non-nil WaitFunc or a non-nil error.
@@ -317,19 +327,24 @@ func (e *Engine) Execute(s *State, file string, script *bufio.Reader, log io.Wri
 		if err != nil {
 			if cmd.want == successRetryOnFailure || cmd.want == failureRetryOnSuccess {
 				// Command wants retries. Retry the whole section
+				numRetries := 0
 				for err != nil {
+					s.FlushLog()
 					select {
 					case <-s.Context().Done():
 						return lineErr(s.Context().Err())
 					case <-time.After(retryInterval):
 					}
+					s.Logf("(command %q failed, retrying...)\n", line)
+					numRetries++
 					for _, cmd := range sectionCmds {
 						impl := e.Cmds[cmd.name]
 						if err = e.runCommand(s, cmd, impl); err != nil {
 							break
 						}
 					}
 				}
+				s.Logf("(command %q succeeded after %d retries)\n", line, numRetries)
 			} else {
 				if stop := (stopError{}); errors.As(err, &stop) {
 					// Since the 'stop' command halts execution of the entire script,
@@ -665,11 +680,42 @@ func (e *Engine) runCommand(s *State, cmd *command, impl Cmd) error {
 		return cmdError(cmd, errors.New("unknown command"))
 	}
 
-	async := impl.Usage().Async
+	usage := impl.Usage()
+
+	async := usage.Async
 	if cmd.background && !async {
 		return cmdError(cmd, errors.New("command cannot be run in background"))
 	}
 
+	// Register and parse the flags. We do this regardless of whether [usage.Flags]
+	// is set in order to handle -h/--help.
+	fs := pflag.NewFlagSet(cmd.name, pflag.ContinueOnError)
+	fs.Usage = func() {} // Don't automatically handle error
+	fs.SetOutput(s.logOut)
+	if usage.Flags != nil {
+		usage.Flags(fs)
+	}
+	if err := fs.Parse(cmd.args); err != nil {
+		if errors.Is(err, pflag.ErrHelp) {
+			out := new(strings.Builder)
+			err = e.ListCmds(out, true, "^"+cmd.name+"$")
+			s.stdout = out.String()
+			if s.stdout != "" {
+				s.Logf("[stdout]\n%s", s.stdout)
+			}
+			return err
+		}
+		if usage.Flags != nil {
+			return cmdError(cmd, err)
+		}
+
+		// [usage.Flags] wasn't given, so ignore the parsing errors as the
+		// command might do its own argument parsing.
+	} else {
+		cmd.args = fs.Args()
+		s.Flags = fs
+	}
+
 	wait, runErr := impl.Run(s, cmd.args...)
 	if wait == nil {
 		if async && runErr == nil {
@@ -788,22 +834,64 @@ func (e *Engine) ListCmds(w io.Writer, verbose bool, regexMatch string) error {
 			suffix = " [&]"
 		}
 
-		_, err := fmt.Fprintf(w, "%s %s%s\n\t%s\n", name, usage.Args, suffix, usage.Summary)
+		flagArgs := ""
+		flagUsages := ""
+		if usage.Flags != nil {
+			fs := pflag.NewFlagSet(name, pflag.ContinueOnError)
+			fs.SetOutput(w)
+			usage.Flags(fs)
+			flagUsages = fs.FlagUsages()
+			shortArgs := []string{}
+			longArgs := []string{}
+			fs.VisitAll(func(flag *pflag.Flag) {
+				if flag.Shorthand != "" {
+					shortArgs = append(shortArgs, flag.Shorthand)
+				}
+				switch flag.Value.Type() {
+				case "bool":
+					longArgs = append(longArgs, fmt.Sprintf("[--%s]", flag.Name))
+				default:
+					longArgs = append(longArgs, fmt.Sprintf("[--%s=%s]", flag.Name, flag.Value.Type()))
+				}
+			})
+			if len(shortArgs) > 0 {
+				flagArgs = fmt.Sprintf("[-%s] ", strings.Join(shortArgs, ""))
+			}
+			flagArgs += strings.Join(longArgs, " ") + " "
+		}
+
+		_, err := fmt.Fprintf(w, "%s %s%s\n\t%s\n", name, flagArgs+usage.Args, suffix, usage.Summary)
 		if err != nil {
 			return err
 		}
 
 		if verbose {
-			if _, err := io.WriteString(w, "\n"); err != nil {
-				return err
-			}
-			for _, line := range usage.Detail {
-				if err := wrapLine(w, line, 60, "\t"); err != nil {
+			if len(usage.Detail) > 0 {
+				if _, err := io.WriteString(w, "\n"); err != nil {
 					return err
 				}
+
+				for _, line := range usage.Detail {
+					if err := wrapLine(w, line, 60, "\t"); err != nil {
+						return err
+					}
+				}
 			}
-			if _, err := io.WriteString(w, "\n"); err != nil {
-				return err
+			if flagUsages != "" {
+				if _, err := io.WriteString(w, "\n"); err != nil {
+					return err
+				}
+				lines := strings.Split(flagUsages, "\n")
+				if len(lines) > 0 {
+					if _, err := fmt.Fprintf(w, "\tFlags:\n"); err != nil {
+						return err
+					}
+					for _, line := range lines {
+						if _, err := fmt.Fprintf(w, "\t%s\n", line); err != nil {
+							return err
+						}
+					}
+				}
 			}
 		}
 	}

script/scripttest/testdata/basic.txt

@@ -22,17 +22,6 @@ echo hello
 
 wait
 
-# Test help in various ways
-help
-help wait
-stdout wait
-help -v wait
-stdout wait
-help unknowncommand
-help ^e
-! stdout wait
-stdout ^exec
-stdout ^exists
-
 -- hello.txt --
 hello world
+

script/scripttest/testdata/help.txt

@@ -0,0 +1,36 @@
+# Test help in various ways
+help
+help wait
+stdout wait
+help -v wait
+stdout wait
+help unknowncommand
+help ^e
+! stdout wait
+stdout ^exec
+stdout ^exists
+
+# Check the 'help' output of the 'stdout' command.
+# The 'cmp' has special handling for 'stdout' and 'stderr'.
+help stdout
+cmp stdout stdout_help.txt
+
+# The -h and --help are intercepted and we show the same
+# output as 'help'
+stdout -h
+cmp stdout stdout_help.txt
+stdout --help
+cmp stdout stdout_help.txt
+
+-- stdout_help.txt --
+stdout [-q] [--count=int] [--quiet] 'pattern'
+	find lines in the stdout buffer that match a pattern
+
+	The command succeeds if at least one match (or the exact
+	count, if given) is found.
+	The -q flag suppresses printing of matches.
+
+	Flags:
+	      --count int   Exact count of matches
+	  -q, --quiet       Suppress printing of matches
+