Add Values() (#167)

* add values template * add no comments flag
abice · Jan 28, 2023 · 1e5efaa · 1e5efaa
1 parent 5354791
commit 1e5efaa
Show file tree

Hide file tree

Showing 12 changed files with 123 additions and 37 deletions.
diff --git a/_example/example.go b/_example/example.go
@@ -1,10 +1,9 @@
-//go:generate ../bin/go-enum --marshal --nocase --flag --names
+//go:generate ../bin/go-enum --marshal --nocase --flag --names --values
 
 package example
 
 // X is doc'ed
-type X struct {
-}
+type X struct{}
 
 // Make x ENUM(Toyota,_,Chevy,_,Ford,_,Tesla,_,Hyundai,_,Nissan,_,Jaguar,_,Audi,_,BMW,_,Mercedes-Benz,_,Volkswagon)
 type Make int32

diff --git a/_example/example_enum.go b/_example/example_enum.go
diff --git a/_example/sql.go b/_example/sql.go
@@ -1,4 +1,4 @@
-//go:generate ../bin/go-enum  --sql --sqlnullstr --sqlnullint --ptr --marshal
+//go:generate ../bin/go-enum  --sql --sqlnullstr --sqlnullint --ptr --marshal --nocomments
 
 package example
 

diff --git a/_example/sql_enum.go b/_example/sql_enum.go
diff --git a/_example/strings_only.go b/_example/strings_only.go
@@ -1,4 +1,4 @@
-//go:generate ../bin/go-enum  --ptr --marshal --flag --nocase --mustparse --sqlnullstr --sql --names
+//go:generate ../bin/go-enum  --ptr --marshal --flag --nocase --mustparse --sqlnullstr --sql --names --values --nocomments
 
 package example
 

diff --git a/_example/strings_only_enum.go b/_example/strings_only_enum.go
diff --git a/generator/enum.tmpl b/generator/enum.tmpl
@@ -16,10 +16,12 @@ import (
 const (
 {{- $enumName := .enum.Name -}}
 {{- $enumType := .enum.Type -}}
+{{- $noComments := .nocomments -}}
 {{- $vars := dict "lastoffset" "0" -}}
 {{ range $rIndex, $value := .enum.Values }}
 	{{- $lastOffset := pluck "lastoffset" $vars | first }}{{ $offset := offset $rIndex $enumType $value }}
-	{{ if eq $value.Name "_"}}// Skipped value.{{else}}// {{$value.PrefixedName}} is a {{$enumName}} of type {{$value.Name}}.{{end}}
+	{{- if $noComments }}{{else}}
+	{{ if eq $value.Name "_"}}// Skipped value.{{else}}// {{$value.PrefixedName}} is a {{$enumName}} of type {{$value.Name}}.{{end}}{{end}}
 	{{- if $value.Comment}}
 	// {{$value.Comment}}
 	{{- end}}
@@ -340,4 +342,15 @@ func {{.enum.Name}}Names() []string {
 }
 {{ end -}}
 
+{{ if .values }}
+
+// {{.enum.Name}}Values returns a list of the values for {{.enum.Name}}
+func {{.enum.Name}}Values() []{{.enum.Name}} {
+    return []{{.enum.Name}}{ {{ range $rIndex, $value := .enum.Values }}{{ if ne $value.Name "_"}}
+		{{$value.PrefixedName}},{{ end }}
+{{- end}}
+    }
+}
+{{ end -}}
+
 {{end}}
diff --git a/generator/enum_string.tmpl b/generator/enum_string.tmpl
@@ -2,9 +2,11 @@
 const (
 {{- $enumName := .enum.Name -}}
 {{- $enumType := .enum.Type -}}
+{{- $noComments := .nocomments -}}
 {{- $vars := dict "lastoffset" "0" -}}
 {{ range $rIndex, $value := .enum.Values }}
-	{{ if eq $value.Name "_"}}// Skipped value.{{else}}// {{$value.PrefixedName}} is a {{$enumName}} of type {{$value.RawName}}.{{end}}
+	{{- if $noComments }}{{else}}
+	{{ if eq $value.Name "_"}}// Skipped value.{{else}}// {{$value.PrefixedName}} is a {{$enumName}} of type {{$value.RawName}}.{{end}}{{end}}
 	{{- if $value.Comment}}
 	// {{$value.Comment}}
 	{{- end}}
@@ -27,6 +29,18 @@ func {{.enum.Name}}Names() []string {
 }
 {{ end -}}
 
+
+{{ if .values }}
+
+// {{.enum.Name}}Values returns a list of the values for {{.enum.Name}}
+func {{.enum.Name}}Values() []{{.enum.Name}} {
+    return []{{.enum.Name}}{ {{ range $rIndex, $value := .enum.Values }}{{ if ne $value.Name "_"}}
+		{{$value.PrefixedName}},{{ end }}
+{{- end}}
+    }
+}
+{{ end -}}
+
 // String implements the Stringer interface.
 func (x {{.enum.Name}}) String() string {
 	return string(x)

diff --git a/generator/generator.go b/generator/generator.go
@@ -25,9 +25,7 @@ const (
 	parseCommentPrefix = `//`
 )
 
-var (
-	replacementNames = map[string]string{}
-)
+var replacementNames = map[string]string{}
 
 // Generator is responsible for generating validation files for the given in a go source file.
 type Generator struct {
@@ -47,13 +45,15 @@ type Generator struct {
 	sqlint            bool
 	flag              bool
 	names             bool
+	values            bool
 	leaveSnakeCase    bool
 	prefix            string
 	sqlNullInt        bool
 	sqlNullStr        bool
 	ptr               bool
 	mustParse         bool
 	forceLower        bool
+	noComments        bool
 }
 
 // Enum holds data for a discovered enum in the parsed source
@@ -154,6 +154,12 @@ func (g *Generator) WithNames() *Generator {
 	return g
 }
 
+// WithValues is used to add Values methods to the enum
+func (g *Generator) WithValues() *Generator {
+	g.values = true
+	return g
+}
+
 // WithoutSnakeToCamel is used to add flag methods to the enum
 func (g *Generator) WithoutSnakeToCamel() *Generator {
 	g.leaveSnakeCase = true
@@ -196,6 +202,12 @@ func (g *Generator) WithForceLower() *Generator {
 	return g
 }
 
+// WithNoComments is used to remove auto generated comments from the enum.
+func (g *Generator) WithNoComments() *Generator {
+	g.noComments = true
+	return g
+}
+
 // ParseAliases is used to add aliases to replace during name sanitization.
 func ParseAliases(aliases []string) error {
 	aliasMap := map[string]string{}
@@ -238,7 +250,6 @@ func (g *Generator) GenerateFromFile(inputFile string) ([]byte, error) {
 		return nil, fmt.Errorf("generate: error parsing input file '%s': %s", inputFile, err)
 	}
 	return g.Generate(f)
-
 }
 
 // Generate does the heavy lifting for the code generation starting from the parsed AST file.
@@ -285,11 +296,13 @@ func (g *Generator) Generate(f *ast.File) ([]byte, error) {
 			"name":       name,
 			"lowercase":  g.lowercaseLookup,
 			"nocase":     g.caseInsensitive,
+			"nocomments": g.noComments,
 			"marshal":    g.marshal,
 			"sql":        g.sql,
 			"sqlint":     g.sqlint,
 			"flag":       g.flag,
 			"names":      g.names,
+			"values":     g.values,
 			"ptr":        g.ptr,
 			"sqlnullint": g.sqlNullInt,
 			"sqlnullstr": g.sqlNullStr,
@@ -343,7 +356,6 @@ func (g *Generator) parseFile(fileName string) (*ast.File, error) {
 
 // parseEnum looks for the ENUM(x,y,z) formatted documentation from the type definition
 func (g *Generator) parseEnum(ts *ast.TypeSpec) (*Enum, error) {
-
 	if ts.Doc == nil {
 		return nil, errors.New("No Doc on Enum")
 	}
@@ -674,7 +686,6 @@ func copyGenDeclCommentsToSpecs(x *ast.GenDecl) {
 			}
 		}
 	}
-
 }
 
 // isTypeSpecEnum checks the comments on the type spec to determine if there is an enum

diff --git a/go.mod b/go.mod
@@ -10,9 +10,9 @@ require (
 	github.com/mattn/goveralls v0.0.11
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.1
-	github.com/urfave/cli/v2 v2.23.7
-	golang.org/x/text v0.5.0
-	golang.org/x/tools v0.4.0
+	github.com/urfave/cli/v2 v2.24.1
+	golang.org/x/text v0.6.0
+	golang.org/x/tools v0.5.0
 )
 
 require (
@@ -34,7 +34,7 @@ require (
 	github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect
 	golang.org/x/crypto v0.1.0 // indirect
 	golang.org/x/mod v0.7.0 // indirect
-	golang.org/x/sys v0.3.0 // indirect
+	golang.org/x/sys v0.4.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -60,8 +60,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/urfave/cli/v2 v2.23.7 h1:YHDQ46s3VghFHFf1DdF+Sh7H4RqhcM+t0TmZRJx4oJY=
-github.com/urfave/cli/v2 v2.23.7/go.mod h1:GHupkWPMM0M/sj1a2b4wUrWBPzazNrIjouW6fmdJLxc=
+github.com/urfave/cli/v2 v2.24.1 h1:/QYYr7g0EhwXEML8jO+8OYt5trPnLHS0p3mrgExJ5NU=
+github.com/urfave/cli/v2 v2.24.1/go.mod h1:GHupkWPMM0M/sj1a2b4wUrWBPzazNrIjouW6fmdJLxc=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/valyala/fasttemplate v1.2.1/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
 github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 h1:bAn7/zixMGCfxrRTfdpNzjtPYqr8smhKouy9mxVdGPU=
@@ -88,18 +88,18 @@ golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211103235746-7861aae1554b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.3.0 h1:w8ZOecv6NaNa/zC8944JTU3vz4u6Lagfk4RPQxv92NQ=
-golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18=
+golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.5.0 h1:OLmvp0KP+FVG99Ct/qFiL/Fhk4zp4QQnZ7b2U+5piUM=
-golang.org/x/text v0.5.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
+golang.org/x/text v0.6.0 h1:3XmdazWV+ubf7QgHSTWeykHOci5oeekaGJBLkrkaw4k=
+golang.org/x/text v0.6.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
-golang.org/x/tools v0.4.0 h1:7mTAgkunk3fr4GAloyyCasadO6h9zSsQZbwvcaIciV4=
-golang.org/x/tools v0.4.0/go.mod h1:UE5sM2OK9E/d67R0ANs2xJizIymRP5gJU295PvKXxjQ=
+golang.org/x/tools v0.5.0 h1:+bSpV5HIeWkuvgaMfI3UmKRThoTA5ODJTUd8T17NO+4=
+golang.org/x/tools v0.5.0/go.mod h1:N+Kgy78s5I24c24dU8OfWNEotWjutIs8SnJvn5IDq+k=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

diff --git a/main.go b/main.go
@@ -30,6 +30,7 @@ type rootT struct {
 	Flag              bool
 	Prefix            string
 	Names             bool
+	Values            bool
 	LeaveSnakeCase    bool
 	SQLNullStr        bool
 	SQLNullInt        bool
@@ -38,6 +39,7 @@ type rootT struct {
 	Aliases           cli.StringSlice
 	MustParse         bool
 	ForceLower        bool
+	NoComments        bool
 }
 
 func main() {
@@ -107,6 +109,11 @@ func main() {
 				Usage:       "Generates a 'Names() []string' function, and adds the possible enum values in the error response during parsing",
 				Destination: &argv.Names,
 			},
+			&cli.BoolFlag{
+				Name:        "values",
+				Usage:       "Generates a 'Values() []{{ENUM}}' function.",
+				Destination: &argv.Values,
+			},
 			&cli.BoolFlag{
 				Name:        "nocamel",
 				Usage:       "Removes the snake_case to CamelCase name changing",
@@ -149,6 +156,11 @@ func main() {
 				Usage:       "Forces a camel cased comment to generate lowercased names.",
 				Destination: &argv.ForceLower,
 			},
+			&cli.BoolFlag{
+				Name:        "nocomments",
+				Usage:       "Removes auto generated comments.  If you add your own comments, these will still be created.",
+				Destination: &argv.NoComments,
+			},
 		},
 		Action: func(ctx *cli.Context) error {
 			if err := generator.ParseAliases(argv.Aliases.Value()); err != nil {
@@ -186,6 +198,9 @@ func main() {
 				if argv.Names {
 					g.WithNames()
 				}
+				if argv.Values {
+					g.WithValues()
+				}
 				if argv.LeaveSnakeCase {
 					g.WithoutSnakeToCamel()
 				}
@@ -207,6 +222,9 @@ func main() {
 				if argv.ForceLower {
 					g.WithForceLower()
 				}
+				if argv.NoComments {
+					g.WithNoComments()
+				}
 				if templates := []string(argv.TemplateFileNames.Value()); len(templates) > 0 {
 					for _, t := range templates {
 						if fn, err := globFilenames(t); err != nil {
@@ -246,7 +264,7 @@ func main() {
 						continue
 					}
 
-					mode := int(0644)
+					mode := int(0o644)
 					err = os.WriteFile(outFilePath, raw, os.FileMode(mode))
 					if err != nil {
 						return fmt.Errorf("failed writing to file %s: %s", color.Cyan(outFilePath), color.Red(err))