support for inline doc comments on optional flags (#549)

support for inline doc comments on optional flags

Author: natefinch

mage/args_test.go

@@ -455,6 +455,73 @@ Usage:
 	}
 }
 
+func TestOptionalArgsFlagDocs(t *testing.T) {
+	stderr := &bytes.Buffer{}
+	stdout := &bytes.Buffer{}
+	inv := Invocation{
+		Dir:    "./testdata/optargs",
+		Stderr: stderr,
+		Stdout: stdout,
+		Help:   true,
+		Args:   []string{"flagdocs"},
+	}
+	code := Invoke(inv)
+	if code != 0 {
+		t.Log("stderr:", stderr)
+		t.Log("stdout:", stdout)
+		t.Fatalf("expected code 0, but got %v", code)
+	}
+	actual := stdout.String()
+	expected := `FlagDocs tests that docs on flags are properly displayed when you run mage -h FlagDocs.
+
+Usage:
+
+	mage flagdocs <name> [<flags>]
+
+Flags:
+
+	-greeting=<string>  the message to append to the name
+	-repeat=<int>       the number of times to repeat
+
+`
+	if actual != expected {
+		t.Fatalf("output is not expected:\ngot:  %q\nwant: %q", actual, expected)
+	}
+}
+
+func TestOptionalArgsSingleFlagDoc(t *testing.T) {
+	stderr := &bytes.Buffer{}
+	stdout := &bytes.Buffer{}
+	inv := Invocation{
+		Dir:    "./testdata/optargs",
+		Stderr: stderr,
+		Stdout: stdout,
+		Help:   true,
+		Args:   []string{"singleflagdoc"},
+	}
+	code := Invoke(inv)
+	if code != 0 {
+		t.Log("stderr:", stderr)
+		t.Log("stdout:", stdout)
+		t.Fatalf("expected code 0, but got %v", code)
+	}
+	actual := stdout.String()
+	expected := `SingleFlagDoc tests that a single documented flag shows the Flags section.
+
+Usage:
+
+	mage singleflagdoc <name> [-greeting=<string>]
+
+Flags:
+
+	-greeting=<string>  the greeting to use
+
+`
+	if actual != expected {
+		t.Fatalf("output is not expected:\ngot:  %q\nwant: %q", actual, expected)
+	}
+}
+
 func TestOptionalArgsCaseInsensitive(t *testing.T) {
 	stderr := &bytes.Buffer{}
 	stdout := &bytes.Buffer{}

mage/template.go

@@ -372,7 +372,9 @@ Options:
 				_fmt.Println({{printf "%q" .Comment}})
 				_fmt.Println()
 				{{end}}
-				_fmt.Print("Usage:\n\n\t{{$.BinaryName}} {{lower .TargetName}}{{range .RequiredArgs}} <{{.Name}}>{{end}}{{range .OptionalArgs}} [-{{.Name}}=<{{.Type}}>]{{end}}\n\n")
+				_fmt.Print("Usage:\n\n\t{{$.BinaryName}} {{lower .TargetName}}{{range .RequiredArgs}} <{{.Name}}>{{end}}{{if .MultipleOptionalArgs}} [<flags>]{{else}}{{range .OptionalArgs}} [-{{.Name}}=<{{.Type}}>]{{end}}{{end}}\n\n")
+				{{if .ShowFlagDocs}}_fmt.Print({{printf "%q" .FlagDocsString}})
+				{{end -}}
 				var aliases []string
 				{{- $name := .Name -}}
 				{{- $recv := .Receiver -}}
@@ -391,7 +393,9 @@ Options:
 				_fmt.Println({{printf "%q" .Comment}})
 				_fmt.Println()
 				{{end}}
-				_fmt.Print("Usage:\n\n\t{{$.BinaryName}} {{lower .TargetName}}{{range .RequiredArgs}} <{{.Name}}>{{end}}{{range .OptionalArgs}} [-{{.Name}}=<{{.Type}}>]{{end}}\n\n")
+				_fmt.Print("Usage:\n\n\t{{$.BinaryName}} {{lower .TargetName}}{{range .RequiredArgs}} <{{.Name}}>{{end}}{{if .MultipleOptionalArgs}} [<flags>]{{else}}{{range .OptionalArgs}} [-{{.Name}}=<{{.Type}}>]{{end}}{{end}}\n\n")
+				{{if .ShowFlagDocs}}_fmt.Print({{printf "%q" .FlagDocsString}})
+				{{end -}}
 				var aliases []string
 				{{- $name := .Name -}}
 				{{- $recv := .Receiver -}}

mage/testdata/optargs/magefile.go

@@ -99,3 +99,27 @@ func Mixed(ctx context.Context, name string, greeting *string, count int) error
 	}
 	return nil
 }
+
+// FlagDocs tests that docs on flags are properly displayed when you
+// run mage -h FlagDocs.
+func FlagDocs(ctx context.Context, name string,
+	greeting *string, // the message to append to the name
+	repeat *int, // the number of times to repeat
+) error {
+	// should show up like this with mage -h FlagDocs below the current usage text:
+	//
+	// -greeting=<string>  the message to append to the name
+	// -repeat=<int>       the number of times to repeat
+	return nil
+}
+
+// SingleFlagDoc tests that a single documented flag shows the Flags section.
+func SingleFlagDoc(name string,
+	greeting *string, // the greeting to use
+) {
+	if greeting != nil {
+		fmt.Printf("%s, %s!\n", *greeting, name)
+	} else {
+		fmt.Printf("Hello, %s!\n", name)
+	}
+}

parse/parse.go

@@ -78,6 +78,7 @@ func (s Functions) Swap(i, j int) {
 type Arg struct {
 	Name, Type string
 	Optional   bool
+	Comment    string
 }
 
 // ID returns user-readable information about where this function is defined.
@@ -149,6 +150,71 @@ func (f Function) HasOptionalArgs() bool {
 	return false
 }
 
+// MultipleOptionalArgs reports whether the function has more than one optional argument.
+func (f Function) MultipleOptionalArgs() bool {
+	n := 0
+	for _, a := range f.Args {
+		if a.Optional {
+			n++
+			if n > 1 {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// ShowFlagDocs reports whether the Flags section should be displayed.
+// This is true when there are multiple optional args (since they are
+// condensed to [<flags>] in the usage line) or when any optional arg
+// has a doc comment.
+func (f Function) ShowFlagDocs() bool {
+	if f.MultipleOptionalArgs() {
+		return true
+	}
+	for _, a := range f.Args {
+		if a.Optional && a.Comment != "" {
+			return true
+		}
+	}
+	return false
+}
+
+// FlagDocsString returns a formatted string documenting optional arguments.
+// It aligns comments to the same column based on the longest flag name.
+func (f Function) FlagDocsString() string {
+	opts := f.OptionalArgs()
+	if len(opts) == 0 {
+		return ""
+	}
+	// Compute flag labels and find max width for alignment
+	type entry struct {
+		label   string
+		comment string
+	}
+	var entries []entry
+	maxLen := 0
+	for _, a := range opts {
+		label := fmt.Sprintf("-%s=<%s>", a.Name, a.Type)
+		if len(label) > maxLen {
+			maxLen = len(label)
+		}
+		entries = append(entries, entry{label: label, comment: a.Comment})
+	}
+
+	var buf strings.Builder
+	_, _ = buf.WriteString("Flags:\n\n")
+	for _, e := range entries {
+		if e.comment != "" {
+			_, _ = fmt.Fprintf(&buf, "\t%-*s  %s\n", maxLen, e.label, e.comment)
+		} else {
+			_, _ = fmt.Fprintf(&buf, "\t%s\n", e.label)
+		}
+	}
+	_, _ = buf.WriteString("\n")
+	return buf.String()
+}
+
 // ExecCode returns code for the template switch to run the target.
 // It wraps each target call to match the func(context.Context) error that
 // runTarget requires.
@@ -413,7 +479,23 @@ func Package(path string, files []string, multiline bool) (*PkgInfo, error) {
 		debug.Printf("found %s tag, using multiline descriptions", multilineTag)
 		multiline = true
 	}
-	p := doc.New(pkg, "./", 0)
+	p := doc.New(pkg, "./", doc.PreserveAST)
+
+	// Build a map from AST fields to their inline comments. We use
+	// ast.NewCommentMap because the Go parser does not populate
+	// ast.Field.Comment for function parameters (only for struct fields).
+	fieldComments := make(map[*ast.Field]string)
+	for _, f := range pkg.Files {
+		cmap := ast.NewCommentMap(fset, f, f.Comments)
+		for node, groups := range cmap {
+			field, ok := node.(*ast.Field)
+			if !ok || len(groups) == 0 {
+				continue
+			}
+			fieldComments[field] = strings.TrimSpace(groups[0].Text())
+		}
+	}
+
 	pi := &PkgInfo{
 		AstPkg:    pkg,
 		DocPkg:    p,
@@ -425,8 +507,8 @@ func Package(path string, files []string, multiline bool) (*PkgInfo, error) {
 		pi.Description = toOneLine(p.Doc)
 	}
 
-	setNamespaces(pi)
-	setFuncs(pi)
+	setNamespaces(pi, fieldComments)
+	setFuncs(pi, fieldComments)
 
 	hasDupes, names := checkDupeTargets(pi)
 	if hasDupes {
@@ -517,29 +599,29 @@ func (s Imports) Swap(i, j int) {
 	s[i], s[j] = s[j], s[i]
 }
 
-func setFuncs(pi *PkgInfo) {
+func setFuncs(pi *PkgInfo, fieldComments map[*ast.Field]string) {
 	for _, f := range pi.DocPkg.Funcs {
 		if f.Recv != "" {
 			debug.Printf("skipping method %s.%s", f.Recv, f.Name)
 			// skip methods
 			continue
 		}
-		fn, ok := funcFromDoc(f, pi.DocPkg.ImportPath, f.Name, pi.Multiline)
+		fn, ok := funcFromDoc(f, pi.DocPkg.ImportPath, f.Name, pi.Multiline, fieldComments)
 		if !ok {
 			continue
 		}
 		pi.Funcs = append(pi.Funcs, fn)
 	}
 }
 
-func setNamespaces(pi *PkgInfo) {
+func setNamespaces(pi *PkgInfo, fieldComments map[*ast.Field]string) {
 	for _, t := range pi.DocPkg.Types {
 		if !isNamespace(t) {
 			continue
 		}
 		debug.Printf("found namespace %s %s", pi.DocPkg.ImportPath, t.Name)
 		for _, f := range t.Methods {
-			fn, ok := funcFromDoc(f, pi.DocPkg.ImportPath, t.Name+"."+f.Name, pi.Multiline)
+			fn, ok := funcFromDoc(f, pi.DocPkg.ImportPath, t.Name+"."+f.Name, pi.Multiline, fieldComments)
 			if !ok {
 				continue
 			}
@@ -549,11 +631,11 @@ func setNamespaces(pi *PkgInfo) {
 	}
 }
 
-func funcFromDoc(f *doc.Func, importpath, funcname string, multiline bool) (*Function, bool) {
+func funcFromDoc(f *doc.Func, importpath, funcname string, multiline bool, fieldComments map[*ast.Field]string) (*Function, bool) {
 	if !ast.IsExported(f.Name) {
 		return nil, false
 	}
-	fn, err := funcType(f.Decl.Type)
+	fn, err := funcType(f.Decl.Type, fieldComments)
 	if err != nil {
 		debug.Printf("skipping invalid method %s %s: %v", importpath, funcname, err)
 		return nil, false
@@ -987,7 +1069,7 @@ func hasErrorReturn(ft *ast.FuncType) (bool, error) {
 	return false, errors.New("EBADRETURNTYPE")
 }
 
-func funcType(ft *ast.FuncType) (*Function, error) {
+func funcType(ft *ast.FuncType, fieldComments map[*ast.Field]string) (*Function, error) {
 	var err error
 	f := &Function{}
 	f.IsContext, err = hasContextParam(ft)
@@ -998,6 +1080,7 @@ func funcType(ft *ast.FuncType) (*Function, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	x := 0
 	if f.IsContext {
 		x++
@@ -1019,9 +1102,10 @@ func funcType(ft *ast.FuncType) (*Function, error) {
 			}
 			return nil, fmt.Errorf("unsupported argument type: %s", t)
 		}
+		comment := fieldComments[param]
 		// support for foo, bar string
 		for _, name := range param.Names {
-			f.Args = append(f.Args, Arg{Name: name.Name, Type: typ, Optional: optional})
+			f.Args = append(f.Args, Arg{Name: name.Name, Type: typ, Optional: optional, Comment: comment})
 		}
 	}
 	return f, nil

parse/parse_test.go

@@ -284,6 +284,14 @@ func TestOptionalArgs(t *testing.T) {
 				{Name: "b", Type: "int", Optional: true},
 			},
 		},
+		{
+			Name: "FlagDocFunc",
+			Args: []Arg{
+				{Name: "name", Type: "string"},
+				{Name: "greeting", Type: "string", Optional: true, Comment: "the greeting message"},
+				{Name: "count", Type: "int", Optional: true, Comment: "how many times"},
+			},
+		},
 		{
 			Name: "OptionalBool",
 			Args: []Arg{

parse/testdata/optargs.go

@@ -16,3 +16,9 @@ func OptionalBool(verbose *bool) {}
 func OptionalDuration(base time.Duration, extra *time.Duration) {}
 
 func AllOptional(a *string, b *int) {}
+
+func FlagDocFunc(name string,
+	greeting *string, // the greeting message
+	count *int, // how many times
+) {
+}