diff --git a/generator/generator.go b/generator/generator.go index cac315a..eaeac05 100644 --- a/generator/generator.go +++ b/generator/generator.go @@ -501,12 +501,14 @@ func getEnumDeclFromComments(comments []*ast.Comment) string { start := strings.Index(line, EnumPrefix) line = line[start+len(EnumPrefix):] } - lineParamLevel := strings.Count(line, "(") - lineParamLevel = lineParamLevel - strings.Count(line, ")") + // we need to ignore anything after the value comment + lineWithoutCommentSuffix, _, _ := strings.Cut(line, parseCommentPrefix) + lineParamLevel := strings.Count(lineWithoutCommentSuffix, "(") + lineParamLevel = lineParamLevel - strings.Count(lineWithoutCommentSuffix, ")") if enumParamLevel+lineParamLevel < 1 { // We've ended, either with more than we need, or with just enough. Now we need to find the end. - for lineIdx, ch := range line { + for lineIdx, ch := range lineWithoutCommentSuffix { if ch == '(' { enumParamLevel = enumParamLevel + 1 continue @@ -516,7 +518,7 @@ func getEnumDeclFromComments(comments []*ast.Comment) string { if enumParamLevel == 0 { // We've found the end of the ENUM() definition, // Cut off the suffix and break out of the loop - line = line[:lineIdx] + line = lineWithoutCommentSuffix[:lineIdx] store = false break } diff --git a/generator/generator_test.go b/generator/generator_test.go index 23f2264..c22b848 100644 --- a/generator/generator_test.go +++ b/generator/generator_test.go @@ -469,6 +469,29 @@ func TestQuotedStrings(t *testing.T) { } } +// TestValueComments tests that we can override the comment of a value +// AND that anything after a value comment (// in //) is ignored, even if that +// would close the enum declaration. +func TestValueComments(t *testing.T) { + input := `package test + // This is a pre-enum comment that needs (to be handled properly) + //ENUM( + // dog, // woof :) + // cat, // meow :3 + //) + type Animal string + ` + g := NewGenerator(WithNoComments()) + + f, err := parser.ParseFile(g.fileSet, "TestRequiredErrors", input, parser.ParseComments) + assert.Nil(t, err, "Error parsing no struct input") + + output, err := g.Generate(f) + assert.Nil(t, err, "Error generating formatted code") + fmt.Println(string(output)) + assert.Contains(t, string(output), "\t// woof :)\n\tAnimalDog Animal = \"dog\"\n\t// meow :3\n\tAnimalCat Animal = \"cat\"") +} + func TestStringWithSingleDoubleQuoteValue(t *testing.T) { input := `package test // ENUM(DoubleQuote='"')