main.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/lindell/string-enumer/pkg/stringenumer"
	"github.com/spf13/pflag"
)

var (
	types      = pflag.StringSliceP("type", "t", nil, "the type name(s), can be multiple, but at least on must be set")
	text       = pflag.BoolP("text", "T", false, "if set, text unmarshaling methods will be generated. Default: false")
	outputPath = pflag.StringP("output", "o", "", "output file name; default is stdout")
)

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of string-enumer:\n")
	fmt.Fprintf(os.Stderr, "\tstring-enumer [flags] --type T --type T2 [directory]\n")
	fmt.Fprintf(os.Stderr, "\tstring-enumer [flags] --type T --type T2 files... # Must be a single package\n")
	fmt.Fprintf(os.Stderr, "For more information, see:\n")
	fmt.Fprintf(os.Stderr, "\thttps://github.com/lindell/string-enumer\n")
	fmt.Fprintf(os.Stderr, "Flags:\n")
	pflag.PrintDefaults()
}

func main() {
	pflag.Usage = Usage
	pflag.Parse()

	if len(*types) == 0 {
		fmt.Fprintln(os.Stderr, "at least one --type needs to be set")
		pflag.Usage()
		os.Exit(2)
	}

	args := pflag.Args()
	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "a directory or file must be defined")
		pflag.Usage()
		os.Exit(2)
	}

	r, err := stringenumer.Generate(
		stringenumer.Paths(args...),
		stringenumer.TypeNames(*types...),
		stringenumer.TextUnmarshaling(*text),
	)
	if err != nil {
		log.Fatalln(err)
	}

	var output io.Writer
	if *outputPath == "" {
		output = os.Stdout
	} else {
		file, err := os.OpenFile(*outputPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		output = file
		defer file.Close()
		if err != nil {
			log.Fatalln(err)
		}
	}

	input := io.MultiReader(generateDontEdit(), r)
	if _, err := io.Copy(output, input); err != nil {
		log.Fatalln(err)
	}
}

func generateDontEdit() io.Reader {
	var buffer bytes.Buffer
	fmt.Fprintf(&buffer, "// Code generated by \"string-enumer %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
	return &buffer
}