Merge pull request #4 from dmarkham/comments-tmpfile

Comments tmpfile
This commit is contained in:
Dan Markham 2019-04-04 16:36:31 -07:00 committed by GitHub
commit d82b2b06c1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 41 additions and 7 deletions

View File

@ -5,7 +5,7 @@
// Enumer is a tool to generate Go code that adds useful methods to Go enums (constants with a specific type).
// It started as a fork of Rob Pikes Stringer tool
//
// Please visit http://github.com/alvaroloes/enumer for a comprehensive documentation
// Please visit http://github.com/dmarkham/enumer for a comprehensive documentation
package main
import (
@ -18,7 +18,6 @@ import (
"go/importer"
"go/token"
"go/types"
"golang.org/x/tools/go/packages"
"io/ioutil"
"log"
"os"
@ -26,9 +25,22 @@ import (
"sort"
"strings"
"golang.org/x/tools/go/packages"
"github.com/pascaldekloe/name"
)
type arrayFlags []string
func (af arrayFlags) String() string {
return strings.Join(af, "")
}
func (af *arrayFlags) Set(value string) error {
*af = append(*af, value)
return nil
}
var (
typeNames = flag.String("type", "", "comma-separated list of type names; must be set")
sql = flag.Bool("sql", false, "if true, the Scanner and Valuer interface will be implemented.")
@ -40,14 +52,20 @@ var (
trimPrefix = flag.String("trimprefix", "", "transform each item name by removing a prefix. Default: \"\"")
)
var comments arrayFlags
func init() {
flag.Var(&comments, "comment", "comments to include in generated code, can repeat. Default: \"\"")
}
// Usage is a replacement usage function for the flags package.
func Usage() {
_, _ = fmt.Fprintf(os.Stderr, "Enumer is a tool to generate Go code that adds useful methods to Go enums (constants with a specific type).")
_, _ = fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
_, _ = fmt.Fprintf(os.Stderr, "\tstringer [flags] -type T [directory]\n")
_, _ = fmt.Fprintf(os.Stderr, "\tstringer [flags] -type T files... # Must be a single package\n")
_, _ = fmt.Fprintf(os.Stderr, "\tEnumer [flags] -type T [directory]\n")
_, _ = fmt.Fprintf(os.Stderr, "\tEnumer [flags] -type T files... # Must be a single package\n")
_, _ = fmt.Fprintf(os.Stderr, "For more information, see:\n")
_, _ = fmt.Fprintf(os.Stderr, "\thttp://godoc.org/golang.org/x/tools/cmd/stringer\n")
_, _ = fmt.Fprintf(os.Stderr, "\thttp://godoc.org/github.com/dmarkham/enumer\n")
_, _ = fmt.Fprintf(os.Stderr, "Flags:\n")
flag.PrintDefaults()
}
@ -89,6 +107,7 @@ func main() {
// Print the header and package clause.
g.Printf("// Code generated by \"enumer %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
g.Printf("\n")
g.Printf("// %s\n", comments.String())
g.Printf("package %s", g.pkg.name)
g.Printf("\n")
g.Printf("import (\n")
@ -109,16 +128,31 @@ func main() {
// Format the output.
src := g.format()
// Write to file.
// Figure out filename to write to
outputName := *output
if outputName == "" {
baseName := fmt.Sprintf("%s_enumer.go", typs[0])
outputName = filepath.Join(dir, strings.ToLower(baseName))
}
err := ioutil.WriteFile(outputName, src, 0644)
// Write to tmpfile first
tmpFile, err := ioutil.TempFile("", fmt.Sprintf("%s_enumer_", typs[0]))
if err != nil {
log.Fatalf("creating temporary file for output: %s", err)
}
_, err = tmpFile.Write(src)
if err != nil {
tmpFile.Close()
os.Remove(tmpFile.Name())
log.Fatalf("writing output: %s", err)
}
tmpFile.Close()
// Rename tmpfile to output file
err = os.Rename(tmpFile.Name(), outputName)
if err != nil {
log.Fatalf("moving tempfile to output file: %s", err)
}
}
// isDirectory reports whether the named file is a directory.