mirror of https://github.com/dmarkham/enumer.git
commit
d82b2b06c1
48
stringer.go
48
stringer.go
|
@ -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 Pike’s 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.
|
||||
|
|
Loading…
Reference in New Issue