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] 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