From be93ed15d012348bbb52cf1b055d156d486c3748 Mon Sep 17 00:00:00 2001
From: Drew Silcock <drew.silcock@stfc.ac.uk>
Date: Thu, 28 Mar 2024 15:13:00 +0000
Subject: [PATCH 1/3] Implement getting name override from comment for enum
 ordered consts.

---
 const.go    | 10 ++++++++++
 packages.go | 14 ++++----------
 schema.go   | 41 +++++++++++++++++++++++++++++++++++++++++
 types.go    | 17 ++---------------
 4 files changed, 57 insertions(+), 25 deletions(-)

diff --git a/const.go b/const.go
index 83755103b..4256dc84a 100644
--- a/const.go
+++ b/const.go
@@ -19,6 +19,16 @@ type ConstVariable struct {
 	Pkg     *PackageDefinitions
 }
 
+func (cv *ConstVariable) VariableName() string {
+	if ignoreNameOverride(cv.Name.Name) {
+		return cv.Name.Name[1:]
+	} else if overriddenName := nameOverride(cv.Comment); overriddenName != "" {
+		return overriddenName
+	}
+
+	return cv.Name.Name
+}
+
 var escapedChars = map[uint8]uint8{
 	'n':  '\n',
 	'r':  '\r',
diff --git a/packages.go b/packages.go
index 466db8629..065cf7d30 100644
--- a/packages.go
+++ b/packages.go
@@ -362,21 +362,15 @@ func (pkgDefs *PackagesDefinitions) collectConstEnums(parsedSchemas map[*TypeSpe
 				typeDef.Enums = make([]EnumValue, 0)
 			}
 
-			name := constVar.Name.Name
+			name := constVar.VariableName()
 			if _, ok = constVar.Value.(ast.Expr); ok {
 				continue
 			}
 
 			enumValue := EnumValue{
-				key:   name,
-				Value: constVar.Value,
-			}
-			if constVar.Comment != nil && len(constVar.Comment.List) > 0 {
-				enumValue.Comment = constVar.Comment.List[0].Text
-				enumValue.Comment = strings.TrimPrefix(enumValue.Comment, "//")
-				enumValue.Comment = strings.TrimPrefix(enumValue.Comment, "/*")
-				enumValue.Comment = strings.TrimSuffix(enumValue.Comment, "*/")
-				enumValue.Comment = strings.TrimSpace(enumValue.Comment)
+				key:     name,
+				Value:   constVar.Value,
+				Comment: commentWithoutNameOverride(constVar.Comment),
 			}
 			typeDef.Enums = append(typeDef.Enums, enumValue)
 		}
diff --git a/schema.go b/schema.go
index b3a5b38c1..a60900564 100644
--- a/schema.go
+++ b/schema.go
@@ -4,6 +4,9 @@ import (
 	"errors"
 	"fmt"
 	"github.com/go-openapi/spec"
+	"go/ast"
+	"regexp"
+	"strings"
 )
 
 const (
@@ -134,6 +137,44 @@ func ignoreNameOverride(name string) bool {
 	return len(name) != 0 && name[0] == IgnoreNameOverridePrefix
 }
 
+var overrideNameRegex = regexp.MustCompile(`(?i)^@name\s+(\S+)`)
+
+func nameOverride(commentGroup *ast.CommentGroup) string {
+	if commentGroup == nil {
+		return ""
+	}
+
+	// get alias from comment '// @name '
+	for _, comment := range commentGroup.List {
+		trimmedComment := strings.TrimSpace(strings.TrimLeft(comment.Text, "/"))
+		texts := overrideNameRegex.FindStringSubmatch(trimmedComment)
+		if len(texts) > 1 {
+			return texts[1]
+		}
+	}
+
+	return ""
+}
+
+func commentWithoutNameOverride(commentGroup *ast.CommentGroup) string {
+	if commentGroup == nil {
+		return ""
+	}
+
+	commentBuilder := strings.Builder{}
+	for _, comment := range commentGroup.List {
+		commentText := comment.Text
+		commentText = strings.TrimPrefix(commentText, "//")
+		commentText = strings.TrimPrefix(commentText, "/*")
+		commentText = strings.TrimSuffix(commentText, "*/")
+		commentText = strings.TrimSpace(commentText)
+		commentText = overrideNameRegex.ReplaceAllString(commentText, "")
+		commentText = strings.TrimSpace(commentText)
+		commentBuilder.WriteString(commentText)
+	}
+	return commentBuilder.String()
+}
+
 // IsComplexSchema whether a schema is complex and should be a ref schema
 func IsComplexSchema(schema *spec.Schema) bool {
 	// a enum type should be complex
diff --git a/types.go b/types.go
index 0076a6b40..c6bd5d221 100644
--- a/types.go
+++ b/types.go
@@ -3,7 +3,6 @@ package swag
 import (
 	"go/ast"
 	"go/token"
-	"regexp"
 	"strings"
 
 	"github.com/go-openapi/spec"
@@ -46,20 +45,8 @@ func (t *TypeSpecDef) Name() string {
 func (t *TypeSpecDef) TypeName() string {
 	if ignoreNameOverride(t.TypeSpec.Name.Name) {
 		return t.TypeSpec.Name.Name[1:]
-	} else if t.TypeSpec.Comment != nil {
-		// get alias from comment '// @name '
-		const regexCaseInsensitive = "(?i)"
-		reTypeName, err := regexp.Compile(regexCaseInsensitive + `^@name\s+(\S+)`)
-		if err != nil {
-			panic(err)
-		}
-		for _, comment := range t.TypeSpec.Comment.List {
-			trimmedComment := strings.TrimSpace(strings.TrimLeft(comment.Text, "/"))
-			texts := reTypeName.FindStringSubmatch(trimmedComment)
-			if len(texts) > 1 {
-				return texts[1]
-			}
-		}
+	} else if overriddenName := nameOverride(t.TypeSpec.Comment); overriddenName != "" {
+		return overriddenName
 	}
 
 	var names []string

From 8f8cbc345b8ab7938930e3bb5e7f80541a3ca1a3 Mon Sep 17 00:00:00 2001
From: Drew Silcock <drew.silcock@stfc.ac.uk>
Date: Thu, 28 Mar 2024 15:13:17 +0000
Subject: [PATCH 2/3] Update README and enum tests for enum variant name
 override feature.

---
 README.md                     | 30 ++++++++++++++++++++++++++++++
 enums_test.go                 | 11 +++++++++++
 testdata/enums/types/model.go |  8 ++++++++
 3 files changed, 49 insertions(+)

diff --git a/README.md b/README.md
index 066ba58ea..5b9e081f9 100644
--- a/README.md
+++ b/README.md
@@ -900,6 +900,36 @@ Make it OR condition
 // @Security OAuth2Application[write, admin] || APIKeyAuth
 ```
 
+### Generate enum types from enum constants
+
+You can generate enums from ordered constants. Each enum variant can have a comment, an override name, or both. This works with both iota-defined and manually defined constants.
+
+```go
+type Difficulty string
+
+const (
+	Easy   Difficulty = "easy" // You can add a comment to the enum variant.
+	Medium Difficulty = "medium" // @name MediumDifficulty
+	Hard   Difficulty = "hard" // @name HardDifficulty You can have a name override and a comment.
+)
+
+type Class int
+
+const (
+	First Class = iota // @name FirstClass
+	Second // Name override and comment rules apply here just as above.
+	Third // @name ThirdClass This one has a name override and a comment.
+)
+
+// There is no need to add `enums:"..."` to the fields, it is automatically generated from the ordered consts.
+type Quiz struct {
+	Difficulty Difficulty
+	Class Class
+	Questions []string
+	Answers []string
+}
+```
+
 
 ### Add a description for enum items
 
diff --git a/enums_test.go b/enums_test.go
index e3456378e..533657447 100644
--- a/enums_test.go
+++ b/enums_test.go
@@ -17,9 +17,11 @@ func TestParseGlobalEnums(t *testing.T) {
 	p := New()
 	err = p.ParseAPI(searchDir, mainAPIFile, defaultParseDepth)
 	assert.NoError(t, err)
+
 	b, err := json.MarshalIndent(p.swagger, "", "    ")
 	assert.NoError(t, err)
 	assert.Equal(t, string(expected), string(b))
+
 	constsPath := "github.com/swaggo/swag/testdata/enums/consts"
 	assert.Equal(t, 64, p.packages.packages[constsPath].ConstTable["uintSize"].Value)
 	assert.Equal(t, int32(62), p.packages.packages[constsPath].ConstTable["maxBase"].Value)
@@ -30,4 +32,13 @@ func TestParseGlobalEnums(t *testing.T) {
 	assert.Equal(t, "aa\nbb\u8888cc", p.packages.packages[constsPath].ConstTable["escapestr"].Value)
 	assert.Equal(t, 1_000_000, p.packages.packages[constsPath].ConstTable["underscored"].Value)
 	assert.Equal(t, 0b10001000, p.packages.packages[constsPath].ConstTable["binaryInteger"].Value)
+
+	typesPath := "github.com/swaggo/swag/testdata/enums/types"
+	difficultyEnums := p.packages.packages[typesPath].TypeDefinitions["Difficulty"].Enums
+	assert.Equal(t, "Easy", difficultyEnums[0].key)
+	assert.Equal(t, "", difficultyEnums[0].Comment)
+	assert.Equal(t, "Medium", difficultyEnums[1].key)
+	assert.Equal(t, "This one also has a comment", difficultyEnums[1].Comment)
+	assert.Equal(t, "DifficultyHard", difficultyEnums[2].key)
+	assert.Equal(t, "This means really hard", difficultyEnums[2].Comment)
 }
diff --git a/testdata/enums/types/model.go b/testdata/enums/types/model.go
index a6cba9cd1..f67d6cae4 100644
--- a/testdata/enums/types/model.go
+++ b/testdata/enums/types/model.go
@@ -63,3 +63,11 @@ type PersonWithArrayEnum struct {
 	Mask  []Mask
 	Type  Type
 }
+
+type Difficulty string
+
+const (
+	DifficultyEasy   Difficulty = "easy"   // @name Easy
+	DifficultyMedium Difficulty = "medium" // @Name Medium This one also has a comment
+	DifficultyHard   Difficulty = "hard"   // This means really hard
+)

From 61ef538a799b415aaad4a474aca44f016a4d844c Mon Sep 17 00:00:00 2001
From: Drew Silcock <drew.silcock@stfc.ac.uk>
Date: Mon, 29 Apr 2024 11:53:48 +0100
Subject: [PATCH 3/3] Update to address comments and linting.

---
 const.go | 5 ++++-
 types.go | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/const.go b/const.go
index 4256dc84a..59c1435bb 100644
--- a/const.go
+++ b/const.go
@@ -19,10 +19,13 @@ type ConstVariable struct {
 	Pkg     *PackageDefinitions
 }
 
+// VariableName gets the bane for this const variable, taking into account comment overrides.
 func (cv *ConstVariable) VariableName() string {
 	if ignoreNameOverride(cv.Name.Name) {
 		return cv.Name.Name[1:]
-	} else if overriddenName := nameOverride(cv.Comment); overriddenName != "" {
+	}
+
+	if overriddenName := nameOverride(cv.Comment); overriddenName != "" {
 		return overriddenName
 	}
 
diff --git a/types.go b/types.go
index c6bd5d221..0c0051c87 100644
--- a/types.go
+++ b/types.go
@@ -45,7 +45,9 @@ func (t *TypeSpecDef) Name() string {
 func (t *TypeSpecDef) TypeName() string {
 	if ignoreNameOverride(t.TypeSpec.Name.Name) {
 		return t.TypeSpec.Name.Name[1:]
-	} else if overriddenName := nameOverride(t.TypeSpec.Comment); overriddenName != "" {
+	}
+
+	if overriddenName := nameOverride(t.TypeSpec.Comment); overriddenName != "" {
 		return overriddenName
 	}