diff --git a/README.md b/README.md index ca5e5ac..830c3ce 100644 --- a/README.md +++ b/README.md @@ -337,6 +337,7 @@ GLOBAL OPTIONS: --forcelower Forces a camel cased comment to generate lowercased names. (default: false) --forceupper Forces a camel cased comment to generate uppercased names. (default: false) --nocomments Removes auto generated comments. If you add your own comments, these will still be created. (default: false) + --comments If you add your own comments, a method for getting comments will be added. (default: false) --buildtag value, -b value [ --buildtag value, -b value ] Adds build tags to a generated enum file. --output-suffix .go Changes the default filename suffix of _enum to something else. .go will be appended to the end of the string no matter what, so that `_test.go` cases can be accommodated --no-iota Disables the use of iota in generated enums. (default: false) diff --git a/generator/enum.tmpl b/generator/enum.tmpl index 8d214a5..1d2b9a5 100644 --- a/generator/enum.tmpl +++ b/generator/enum.tmpl @@ -79,6 +79,18 @@ func {{.parseName}}{{.enum.Name}}(name string) ({{.enum.Name}}, error) { } {{- end }} +{{ if .comments }} +var _{{.enum.Name}}Comments = {{ unmapifyComment .enum .lowercase }} + +// {{.enum.Name}}Comment returns the comment associated with the enum value name. +func {{.enum.Name}}Comment(name string) string { + if x, ok := _{{.enum.Name}}Comments[name]; ok { + return x + } + return "" +} +{{ end }} + {{ if .mustparse }} // MustParse{{.enum.Name}} converts a string to a {{.enum.Name}}, and panics if is not valid. func MustParse{{.enum.Name}}(name string) {{.enum.Name}} { diff --git a/generator/enum_string.tmpl b/generator/enum_string.tmpl index ee929d3..7d40dd6 100644 --- a/generator/enum_string.tmpl +++ b/generator/enum_string.tmpl @@ -76,6 +76,18 @@ func {{.parseName}}{{.enum.Name}}(name string) ({{.enum.Name}}, error) { } {{- end }} +{{ if .comments }} +var _{{.enum.Name}}Comments = {{ unmapifyCommentStringEnum .enum .forceupper }} + +// {{.enum.Name}}Comment returns the comment associated with the enum value name. +func {{.enum.Name}}Comment(name string) string { + if x, ok := _{{.enum.Name}}Comments[name]; ok { + return x + } + return "" +} +{{ end }} + {{ if .mustparse }} // MustParse{{.enum.Name}} converts a string to a {{.enum.Name}}, and panics if is not valid. func MustParse{{.enum.Name}}(name string) {{.enum.Name}} { diff --git a/generator/generator.go b/generator/generator.go index cac315a..b99eea0 100644 --- a/generator/generator.go +++ b/generator/generator.go @@ -90,6 +90,8 @@ func NewGeneratorWithConfig(config GeneratorConfig) *Generator { funcs["stringify"] = Stringify funcs["mapify"] = Mapify funcs["unmapify"] = Unmapify + funcs["unmapifyCommentStringEnum"] = UnmapifyCommentStringEnum + funcs["unmapifyComment"] = UnmapifyComment funcs["namify"] = Namify funcs["offset"] = Offset funcs["quote"] = strconv.Quote @@ -212,6 +214,7 @@ func (g *Generator) Generate(f *ast.File) ([]byte, error) { "lowercase": g.LowercaseLookup, "nocase": g.CaseInsensitive, "nocomments": g.NoComments, + "comments": g.Comments, "noIota": g.NoIota, "marshal": g.Marshal, "sql": g.SQL, diff --git a/generator/generator_test.go b/generator/generator_test.go index 23f2264..8377c5d 100644 --- a/generator/generator_test.go +++ b/generator/generator_test.go @@ -989,3 +989,61 @@ type Greek string assert.Contains(t, outputStr, "var ErrInvalidGreek") assert.Contains(t, outputStr, "lookupSqlIntGreek") } + +// TestWithComments +func TestWithComments(t *testing.T) { + input := `package test +// ENUM( +// +// disabled // This User has disabled +// banned // This User has been banned +// active // This User is active +// +// ) +type UserStatus int +` + g := NewGenerator(WithComments()) + f, err := parser.ParseFile(g.fileSet, "test.go", input, parser.ParseComments) + require.NoError(t, err) + + output, err := g.Generate(f) + require.NoError(t, err) + require.NotNil(t, output) + + outputStr := string(output) + + // Should contain error variable because lookupSqlInt and Value use it + assert.Contains(t, outputStr, "var _UserStatusComments = map[string]string{") + assert.Contains(t, outputStr, `_UserStatusName[0:8]: UserStatusDisabled,`) + assert.Contains(t, outputStr, `_UserStatusName[8:14]: UserStatusBanned,`) + assert.Contains(t, outputStr, `_UserStatusName[14:20]: UserStatusActive,`) +} + +// TestStringEnumWithComments +func TestStringEnumWithComments(t *testing.T) { + input := `package test +// ENUM( +// +// disabled // This User has disabled +// banned // This User has been banned +// active // This User is active +// +// ) +type UserStatus string +` + g := NewGenerator(WithComments()) + f, err := parser.ParseFile(g.fileSet, "test.go", input, parser.ParseComments) + require.NoError(t, err) + + output, err := g.Generate(f) + require.NoError(t, err) + require.NotNil(t, output) + + outputStr := string(output) + + // Should contain error variable because lookupSqlInt and Value use it + assert.Contains(t, outputStr, "var _UserStatusComments = map[string]string{") + assert.Contains(t, outputStr, `"disabled": UserStatusDisabled,`) + assert.Contains(t, outputStr, `"banned": UserStatusBanned,`) + assert.Contains(t, outputStr, `"active": UserStatusActive,`) +} diff --git a/generator/options.go b/generator/options.go index df08594..4bb9ffc 100644 --- a/generator/options.go +++ b/generator/options.go @@ -22,6 +22,7 @@ type GeneratorConfig struct { ForceLower bool `json:"force_lower"` ForceUpper bool `json:"force_upper"` NoComments bool `json:"no_comments"` + Comments bool `json:"comments"` NoParse bool `json:"no_parse"` BuildTags []string `json:"build_tags"` ReplacementNames map[string]string `json:"replacement_names"` @@ -180,6 +181,13 @@ func WithNoComments() Option { } } +// WithComments is used to add method for getting comments. +func WithComments() Option { + return func(g *GeneratorConfig) { + g.Comments = true + } +} + // WithBuildTags will add build tags to the generated file. func WithBuildTags(tags ...string) Option { return func(g *GeneratorConfig) { diff --git a/generator/template_funcs.go b/generator/template_funcs.go index c4549dd..7609adb 100644 --- a/generator/template_funcs.go +++ b/generator/template_funcs.go @@ -61,6 +61,54 @@ func Unmapify(e Enum, lowercase bool) (ret string, err error) { return } +// UnmapifyComment returns a map that is all of the indexes for a string value lookup +func UnmapifyComment(e Enum, lowercase bool) (ret string, err error) { + if e.Type == "string" { + return UnmapifyStringEnum(e, lowercase) + } + strName := fmt.Sprintf(`_%sName`, e.Name) + ret = "map[string]string{\n" + index := 0 + for _, val := range e.Values { + if val.Name != skipHolder { + nextIndex := index + len(val.Name) + ret = fmt.Sprintf("%s%s[%d:%d]: \"%s\",\n", ret, strName, index, nextIndex, val.Comment) + if lowercase { + ret = fmt.Sprintf("%sstrings.ToLower(%s[%d:%d]): \"%s\",\n", ret, strName, index, nextIndex, val.Comment) + } + index = nextIndex + } + } + ret = ret + `}` + return +} + +// UnmapifyCommentStringEnum returns a map that is all of the indexes for a string value lookup +func UnmapifyCommentStringEnum(e Enum, lowercase bool) (ret string, err error) { + var builder strings.Builder + _, err = builder.WriteString("map[string]string{\n") + if err != nil { + return + } + for _, val := range e.Values { + if val.Name != skipHolder { + _, err = builder.WriteString(fmt.Sprintf("%q:\"%s\",\n", val.ValueStr, val.Comment)) + if err != nil { + return + } + if lowercase && strings.ToLower(val.ValueStr) != val.ValueStr { + _, err = builder.WriteString(fmt.Sprintf("%q:\"%s\",\n", strings.ToLower(val.ValueStr), val.Comment)) + if err != nil { + return + } + } + } + } + builder.WriteByte('}') + ret = builder.String() + return +} + // Unmapify returns a map that is all of the indexes for a string value lookup func UnmapifyStringEnum(e Enum, lowercase bool) (ret string, err error) { var builder strings.Builder diff --git a/main.go b/main.go index a149bc4..6276f9c 100644 --- a/main.go +++ b/main.go @@ -47,6 +47,7 @@ type rootT struct { ForceLower bool ForceUpper bool NoComments bool + Comments bool NoParse bool OutputSuffix string } @@ -208,6 +209,11 @@ func main() { Usage: "Removes auto generated comments. If you add your own comments, these will still be created.", Destination: &argv.NoComments, }, + &cli.BoolFlag{ + Name: "comments", + Usage: "If you add your own comments, a method for getting comments will be added.", + Destination: &argv.Comments, + }, &cli.BoolFlag{ Name: "noparse", Usage: "Prevents generating the Parse method, or generates it as unexported if other methods depend on it.", @@ -280,6 +286,7 @@ func main() { ForceLower: argv.ForceLower, ForceUpper: argv.ForceUpper, NoComments: argv.NoComments, + Comments: argv.Comments, NoParse: argv.NoParse, BuildTags: argv.BuildTags.Value(), ReplacementNames: aliases,