generator.go

package flagforge

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"strings"
	"text/template"
	"time"
)

const flagTemplate = `
// Code generated by go generate; DO NOT EDIT.
package {{ .Pkg }}

import (
	"errors"
	"flag"
	"fmt"
	"os"
    "strings"
	"time"
)

// {{ .ConfigType }} represents all configuration options.
type {{ .ConfigType }} struct {
{{- range .Args }}
	// {{ .ShortHelp }}
	{{ .Name }} {{ .Type }}
{{- end }}
{{- range .Flags }}
	// {{ .ShortHelp }}
	{{- if eq .Type "filepath" }}
	{{ .Name }} string ` + "`filepath:\"true\"`" + `
	{{- else }}
	{{ .Name }} {{ .Type }}
	{{- end }}
{{- end }}
}

// Forge sets up and parses command-line flags.
func Forge(arguments []string) (*flag.FlagSet, *{{ .ConfigType }}, error) {
	config := &{{ .ConfigType }}{}
	fs := flag.NewFlagSet("{{ .FSName }}", flag.{{ .FSErrorHandling }})
{{- range $index, $element := .Args }}
	if len(arguments) <= {{ $index }} {
		return nil, nil, fmtError("missing required argument: {{ $element.Name }}")
	}
{{- end }}
{{- range .Flags }}
	{{- if or (eq .Type "string") (eq .Type "filepath") }}
	fs.StringVar(&config.{{ .Name }}, "{{ .CLI }}", "{{ .Default }}", "{{ .ShortHelp }}")
	{{- else if eq .Type "bool" }}
	fs.BoolVar(&config.{{ .Name }}, "{{ .CLI }}", {{ .Default }}, "{{ .ShortHelp }}")
	{{- else if eq .Type "int" }}
	fs.IntVar(&config.{{ .Name }}, "{{ .CLI }}", {{ .Default }}, "{{ .ShortHelp }}")
	{{- else if eq .Type "uint64" }}
	fs.Uint64Var(&config.{{ .Name }}, "{{ .CLI }}", {{ .Default }}, "{{ .ShortHelp }}")
	{{- else if eq .Type "int64" }}
	fs.Int64Var(&config.{{ .Name }}, "{{ .CLI }}", {{ .Default }}, "{{ .ShortHelp }}")
	{{- else if eq .Type "time.Duration" }}
	fs.DurationVar(&config.{{ .Name }}, "{{ .CLI }}", mustParseDuration("{{ .Default }}"), "{{ .ShortHelp }}")
	{{- else if eq .Type "[]string" }}
	var tmp{{ .Name }} string
	fs.StringVar(&tmp{{ .Name }}, "{{ .CLI }}", "{{ .Default }}", "{{ .ShortHelp }}")
	{{- end }}
{{- end }}
{{- if .FSUsage }}
	fs.Usage = func() {
		usage("{{ .FSUsage }}")
		fs.PrintDefaults()
	}
{{- end }}
    if err := fs.Parse(arguments); err != nil {
	    return nil, nil, err
    }
{{- range $index, $element := .Args }}
	{{- if eq .Type "string" }}
	    config.{{ .Name }} = fs.Arg({{ $index }})
	{{- end }}
{{- end }}
{{- range $index, $element := .Flags }}
	{{- if eq .Type "[]string" }}
	    config.{{ .Name }} = splitString(tmp{{ .Name }}, "{{ .Delimiter }}")
	{{- end }}
{{- end }}
	return fs, config, nil
}

func mustParseDuration(d string) time.Duration {
	td, err := time.ParseDuration(d)
	if err != nil {
		panic(err)
	}
	return td
}

func splitString(s, sep string) []string {
	if s == "" {
		return nil
	}
	return strings.Split(s, sep)
}

func fmtError(msg string) error {
	return errors.New(msg)
}

func usage(msg string) {
	fmt.Fprintf(os.Stderr, "%s", msg)
}
`

// htmlSectionTemplate renders a single section: an optional Markdown heading
// followed by a table of that section's flags. The output is a fragment, not a
// complete HTML document, so that it can be embedded directly in a page which
// supplies its own styling -- the table carries the "rq-flags" class for that
// purpose. Headings are Markdown rather than <h2> so that a static site
// generator treats them as real headings, and so gives them anchors and a place
// in the page's table of contents.
const htmlSectionTemplate = `{{ if .Name }}## {{ .Name }}

{{ end }}<table class="rq-flags">
	<tr>
		<th class="col-cli">Flag</th>
		<th class="col-usage">Usage</th>
	</tr>
	{{- range .Flags }}
	<tr>
		<td><code>-{{ .CLI | html }}</code></td>
		<td>{{ .ShortHelp | html }}.
		{{- if .LongHelp }}
		    <br><br>{{ .LongHelp | html }}
		{{- end }}</td>
	</tr>
	{{- end }}
</table>
`

// Format represents the output format of the generator.
type Format int

const (
	Go Format = iota
	Markdown
	HTML
)

// String returns the string representation of the format.
func (f Format) String() string {
	switch f {
	case Go:
		return "Go"
	case Markdown:
		return "Markdown"
	case HTML:
		return "HTML"
	default:
		return "Unknown"
	}
}

// visibleFlags returns a copy of flags with hidden entries removed.
func visibleFlags(flags []Flag) []Flag {
	var out []Flag
	for _, f := range flags {
		if !f.Hide {
			out = append(out, f)
		}
	}
	return out
}

// Generator represents a flag, HTML, or Markdown generator.
type Generator struct {
	pkg            string
	configTypeName string

	flagSetUsage         string
	flagSetName          string
	flagSetErrorHandling string

	args  []Argument
	flags []Flag
}

// NewGenerator creates a new generator with the given package name, name, and
// path to the TOML configuration file.
func NewGenerator(cfg *ParsedConfig) (*Generator, error) {
	return &Generator{
		pkg:                  cfg.GoConfig.Package,
		configTypeName:       cfg.GoConfig.ConfigTypeName,
		flagSetUsage:         cfg.GoConfig.FlagSetUsage,
		flagSetName:          cfg.GoConfig.FlagSetName,
		flagSetErrorHandling: cfg.GoConfig.FlagErrorHandling,
		args:                 cfg.Arguments,
		flags:                cfg.Flags,
	}, nil
}

// Execute generates the output in the given format and writes it to the given
// writer.
func (g *Generator) Execute(f Format, w io.Writer) error {
	switch f {
	case Go:
		return g.doGo(w)
	case Markdown:
		return g.doMarkdown(w)
	case HTML:
		return g.doHTML(w)
	default:
		return fmt.Errorf("unsupported format: %s", f)
	}
}

func (g *Generator) doGo(w io.Writer) error {
	// Parse the template.
	tmpl, err := template.New("flags").Parse(flagTemplate)
	if err != nil {
		return fmt.Errorf("failed to parse template: %w", err)
	}

	// Perform some checks of the flags.
	for i, flag := range g.flags {
		if flag.Type == "time.Duration" {
			if flag.Default == nil {
				g.flags[i].Default = 0
			} else {
				s, ok := flag.Default.(string)
				if !ok {
					return fmt.Errorf("time.Duration flag %s has non-string default", flag.Name)
				}
				if _, err := time.ParseDuration(s); err != nil {
					return fmt.Errorf("time.Duration flag %s has invalid default: %v", flag.Name, err)
				}
			}
		}
		if flag.Type == "[]string" {
			if flag.Delimiter == "" {
				g.flags[i].Delimiter = ","
			}
			if flag.Default == nil {
				g.flags[i].Default = ""
			}
		}
	}

	// Execute the template with the flags data.
	var output bytes.Buffer
	if err := tmpl.Execute(&output, struct {
		Pkg             string
		FSUsage         string
		FSName          string
		FSErrorHandling string
		ConfigType      string
		Args            []Argument
		Flags           []Flag
	}{
		Pkg:             g.pkg,
		FSUsage:         g.flagSetUsage,
		FSName:          g.flagSetName,
		FSErrorHandling: g.flagSetErrorHandling,
		ConfigType:      g.configTypeName,
		Args:            g.args,
		Flags:           visibleFlags(g.flags),
	}); err != nil {
		return fmt.Errorf("failed to execute template: %w", err)
	}

	// Format the Go source.
	formatted, err := format.Source(output.Bytes())
	if err != nil {
		return fmt.Errorf("failed to format source: %w", err)
	}

	// Write the output to flags.go.
	_, err = w.Write(formatted)
	return err
}

func (g *Generator) doMarkdown(w io.Writer) error {
	sections, err := groupBySection(visibleFlags(g.flags))
	if err != nil {
		return err
	}

	for i, section := range sections {
		builder := strings.Builder{}
		if i > 0 {
			builder.WriteString("\n")
		}
		if section.Name != "" {
			builder.WriteString(fmt.Sprintf("## %s\n\n", section.Name))
		}

		// Write the markdown table header.
		builder.WriteString("| Flag | Usage |\n|-|-|\n")

		// Write each flag as a row in the table.
		for _, flag := range section.Flags {
			builder.WriteString("|")
			builder.WriteString(escapeMarkdown(flag.CLI))
			builder.WriteString("|")
			builder.WriteString(escapeMarkdown(flag.ShortHelp))
			if flag.Default != nil {
				if !strings.HasSuffix(flag.ShortHelp, ".") {
					builder.WriteString(".")
				}
				builder.WriteString(fmt.Sprintf(" %s", escapeMarkdown(flag.LongHelp)))
			}
			builder.WriteString("|\n")
		}
		if _, err := w.Write([]byte(builder.String())); err != nil {
			return err
		}
	}
	return nil
}

func (g *Generator) doHTML(w io.Writer) error {
	sections, err := groupBySection(visibleFlags(g.flags))
	if err != nil {
		return err
	}

	// Parse the template.
	tmpl, err := template.New("htmlTable").Funcs(template.FuncMap{
		"html": func(s string) string {
			return template.HTMLEscapeString(s)
		},
	}).Parse(htmlSectionTemplate)
	if err != nil {
		return fmt.Errorf("failed to parse HTML template: %w", err)
	}

	// Execute the template once per section, separating each from the last with
	// a blank line.
	var output bytes.Buffer
	for i, section := range sections {
		if i > 0 {
			output.WriteString("\n")
		}
		if err := tmpl.Execute(&output, section); err != nil {
			return fmt.Errorf("failed to execute HTML template: %w", err)
		}
	}

	if _, err := w.Write(output.Bytes()); err != nil {
		return fmt.Errorf("failed to write HTML: %w", err)
	}
	return nil
}

// section is a named group of flags, used by the documentation generators.
type section struct {
	Name  string
	Flags []Flag
}

// groupBySection groups flags by their section key, ordering the sections by
// first appearance in the configuration file. If no flag declares a section a
// single anonymous section holding every flag is returned, so that output for
// configuration files which don't use sections is unchanged.
//
// Declaring a section on some flags but not others is an error. Silently
// sweeping the remainder into a catch-all group would mean that every flag
// added from then on would quietly land there, which is exactly the drift
// sections exist to prevent.
func groupBySection(flags []Flag) ([]section, error) {
	var unassigned []string
	assigned := false
	for _, flag := range flags {
		if flag.Section == "" {
			unassigned = append(unassigned, flag.CLI)
		} else {
			assigned = true
		}
	}
	if !assigned {
		return []section{{Flags: flags}}, nil
	}
	if len(unassigned) > 0 {
		return nil, fmt.Errorf("some flags declare a section but these do not: %s",
			strings.Join(unassigned, ", "))
	}

	var sections []section
	indexes := make(map[string]int)
	for _, flag := range flags {
		i, ok := indexes[flag.Section]
		if !ok {
			i = len(sections)
			indexes[flag.Section] = i
			sections = append(sections, section{Name: flag.Section})
		}
		sections[i].Flags = append(sections[i].Flags, flag)
	}
	return sections, nil
}

// escapeMarkdown escapes markdown special characters.
func escapeMarkdown(text string) string {
	text = strings.ReplaceAll(text, "|", "\\|")
	text = strings.ReplaceAll(text, "\n", "<br>")
	return text
}