From 507c359997202dc2e7888839f478c58ab9b23cad Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Sat, 6 Apr 2019 15:07:37 -0400
Subject: [PATCH] mcfg: modify WithCLITail to allow modifying the Usage
 description

---
 mcfg/cli.go      | 41 +++++++++++++++++++++++++++++++----------
 mcfg/cli_test.go | 36 ++++++++++++++++++++++++++++++------
 2 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/mcfg/cli.go b/mcfg/cli.go
index 7eb9501..62e45da 100644
--- a/mcfg/cli.go
+++ b/mcfg/cli.go
@@ -16,10 +16,15 @@ import (
 type cliKey int
 
 const (
-	cliKeyTailPtr cliKey = iota
+	cliKeyTail cliKey = iota
 	cliKeySubCmdM
 )
 
+type cliTail struct {
+	dst   *[]string
+	descr string
+}
+
 // WithCLITail returns a Context which modifies the behavior of SourceCLI's
 // Parse, if SourceCLI is used with that Context at all. Normally when SourceCLI
 // encounters an unexpected Arg it will immediately return an error. This
@@ -27,21 +32,33 @@ const (
 // and all subsequent Args (i.e. the tail) should be set to the returned
 // []string value.
 //
-// If multiple WithCLITail calls are used then only the latest returned pointer
-// will be filled.
-func WithCLITail(ctx context.Context) (context.Context, *[]string) {
+// The descr (optional) will be appended to the "Usage" line which is printed
+// with the help document when "-h" is passed in.
+func WithCLITail(ctx context.Context, descr string) (context.Context, *[]string) {
+	if ctx.Value(cliKeyTail) != nil {
+		panic("WithCLITail already called in this Context")
+	}
 	tailPtr := new([]string)
-	return context.WithValue(ctx, cliKeyTailPtr, tailPtr), tailPtr
+	ctx = context.WithValue(ctx, cliKeyTail, cliTail{
+		dst:   tailPtr,
+		descr: descr,
+	})
+	return ctx, tailPtr
 }
 
 func populateCLITail(ctx context.Context, tail []string) bool {
-	tailPtr, ok := ctx.Value(cliKeyTailPtr).(*[]string)
+	ct, ok := ctx.Value(cliKeyTail).(cliTail)
 	if ok {
-		*tailPtr = tail
+		*ct.dst = tail
 	}
 	return ok
 }
 
+func getCLITailDescr(ctx context.Context) string {
+	ct, _ := ctx.Value(cliKeyTail).(cliTail)
+	return ct.descr
+}
+
 type subCmd struct {
 	name, descr string
 	flag        *bool
@@ -133,10 +150,9 @@ func (cli *SourceCLI) parse(
 	if err != nil {
 		return nil, nil, err
 	}
-	subCmdM, _ := ctx.Value(cliKeySubCmdM).(map[string]subCmd)
 
 	printHelpAndExit := func() {
-		cli.printHelp(os.Stderr, subCmdPrefix, subCmdM, pM)
+		cli.printHelp(ctx, os.Stderr, subCmdPrefix, pM)
 		os.Stderr.Sync()
 		os.Exit(1)
 	}
@@ -145,6 +161,7 @@ func (cli *SourceCLI) parse(
 	// of them should have been given, in which case send the Context through
 	// the callback to obtain a new one (which presumably has further config
 	// options the previous didn't) and call parse again.
+	subCmdM, _ := ctx.Value(cliKeySubCmdM).(map[string]subCmd)
 	if len(subCmdM) > 0 {
 		subCmd, args, ok := cli.getSubCmd(subCmdM, args)
 		if !ok {
@@ -256,9 +273,9 @@ func (cli *SourceCLI) cliParams(params []Param) (map[string]Param, error) {
 }
 
 func (cli *SourceCLI) printHelp(
+	ctx context.Context,
 	w io.Writer,
 	subCmdPrefix []string,
-	subCmdM map[string]subCmd,
 	pM map[string]Param,
 ) {
 	type pEntry struct {
@@ -297,6 +314,7 @@ func (cli *SourceCLI) printHelp(
 		subCmd
 	}
 
+	subCmdM, _ := ctx.Value(cliKeySubCmdM).(map[string]subCmd)
 	subCmdA := make([]subCmdEntry, 0, len(subCmdM))
 	for name, subCmd := range subCmdM {
 		subCmdA = append(subCmdA, subCmdEntry{name: name, subCmd: subCmd})
@@ -316,6 +334,9 @@ func (cli *SourceCLI) printHelp(
 	if len(pA) > 0 {
 		fmt.Fprint(w, " [options]")
 	}
+	if descr := getCLITailDescr(ctx); descr != "" {
+		fmt.Fprintf(w, " %s", descr)
+	}
 	fmt.Fprint(w, "\n\n")
 
 	if len(subCmdA) > 0 {
diff --git a/mcfg/cli_test.go b/mcfg/cli_test.go
index c713e4b..800ca88 100644
--- a/mcfg/cli_test.go
+++ b/mcfg/cli_test.go
@@ -22,8 +22,7 @@ func TestSourceCLIHelp(t *T) {
 		src := &SourceCLI{}
 		pM, err := src.cliParams(CollectParams(ctx))
 		require.NoError(t, err)
-		subCmdM, _ := ctx.Value(cliKeySubCmdM).(map[string]subCmd)
-		src.printHelp(buf, subCmdPrefix, subCmdM, pM)
+		src.printHelp(ctx, buf, subCmdPrefix, pM)
 
 		out := buf.String()
 		ok := regexp.MustCompile(exp).MatchString(out)
@@ -106,6 +105,31 @@ Options:
 	--foo \(Default: 5\)
 		Test int param.
 
+$`)
+
+	ctx, _ = WithCLITail(ctx, "[arg...]")
+	assertHelp(ctx, nil, `^Usage: \S+ <sub-command> \[options\] \[arg\.\.\.\]
+
+Sub-commands:
+
+	first	First sub-command
+	second	Second sub-command
+
+Options:
+
+	--baz2 \(Required\)
+
+	--baz3 \(Required\)
+
+	--bar \(Flag\)
+		Test bool param.
+
+	--baz \(Default: "baz"\)
+		Test string param.
+
+	--foo \(Default: 5\)
+		Test int param.
+
 $`)
 }
 
@@ -204,7 +228,7 @@ func TestWithCLITail(t *T) {
 	}
 
 	for _, tc := range cases {
-		ctx, tail := WithCLITail(ctx)
+		ctx, tail := WithCLITail(ctx, "foo")
 		_, err := Populate(ctx, &SourceCLI{Args: tc.args})
 		massert.Require(t, massert.Comment(massert.All(
 			massert.Nil(err),
@@ -216,15 +240,15 @@ func TestWithCLITail(t *T) {
 func ExampleWithCLITail() {
 	ctx := context.Background()
 	ctx, foo := WithInt(ctx, "foo", 1, "Description of foo.")
-	ctx, tail := WithCLITail(ctx)
+	ctx, tail := WithCLITail(ctx, "[arg...]")
 	ctx, bar := WithString(ctx, "bar", "defaultVal", "Description of bar.")
 
 	_, err := Populate(ctx, &SourceCLI{
-		Args: []string{"--foo=100", "BADARG", "--bar", "BAR"},
+		Args: []string{"--foo=100", "arg1", "arg2", "arg3"},
 	})
 
 	fmt.Printf("err:%v foo:%v bar:%v tail:%#v\n", err, *foo, *bar, *tail)
-	// Output: err:<nil> foo:100 bar:defaultVal tail:[]string{"BADARG", "--bar", "BAR"}
+	// Output: err:<nil> foo:100 bar:defaultVal tail:[]string{"arg1", "arg2", "arg3"}
 }
 
 func TestWithCLISubCommand(t *T) {