feat: adjust fiber route generation

pkg/swag/formatter.go

@@ -0,0 +1,187 @@
+package swag
+
+import (
+	"bytes"
+	"fmt"
+	"go/ast"
+	goparser "go/parser"
+	"go/token"
+	"log"
+	"os"
+	"regexp"
+	"sort"
+	"strings"
+	"text/tabwriter"
+
+	"golang.org/x/tools/imports"
+)
+
+// Check of @Param @Success @Failure @Response @Header
+var specialTagForSplit = map[string]bool{
+	paramAttr:    true,
+	successAttr:  true,
+	failureAttr:  true,
+	responseAttr: true,
+	headerAttr:   true,
+}
+
+var skipChar = map[byte]byte{
+	'"': '"',
+	'(': ')',
+	'{': '}',
+	'[': ']',
+}
+
+// Formatter implements a formatter for Go source files.
+type Formatter struct {
+	// debugging output goes here
+	debug Debugger
+}
+
+// NewFormatter create a new formatter instance.
+func NewFormatter() *Formatter {
+	formatter := &Formatter{
+		debug: log.New(os.Stdout, "", log.LstdFlags),
+	}
+	return formatter
+}
+
+// Format formats swag comments in contents. It uses fileName to report errors
+// that happen during parsing of contents.
+func (f *Formatter) Format(fileName string, contents []byte) ([]byte, error) {
+	fileSet := token.NewFileSet()
+	ast, err := goparser.ParseFile(fileSet, fileName, contents, goparser.ParseComments)
+	if err != nil {
+		return nil, err
+	}
+
+	// Formatting changes are described as an edit list of byte range
+	// replacements. We make these content-level edits directly rather than
+	// changing the AST nodes and writing those out (via [go/printer] or
+	// [go/format]) so that we only change the formatting of Swag attribute
+	// comments. This won't touch the formatting of any other comments, or of
+	// functions, etc.
+	maxEdits := 0
+	for _, comment := range ast.Comments {
+		maxEdits += len(comment.List)
+	}
+	edits := make(edits, 0, maxEdits)
+
+	for _, comment := range ast.Comments {
+		formatFuncDoc(fileSet, comment.List, &edits)
+	}
+	formatted, err := imports.Process(fileName, edits.apply(contents), nil)
+	if err != nil {
+		return nil, err
+	}
+	return formatted, nil
+}
+
+type edit struct {
+	begin       int
+	end         int
+	replacement []byte
+}
+
+type edits []edit
+
+func (edits edits) apply(contents []byte) []byte {
+	// Apply the edits with the highest offset first, so that earlier edits
+	// don't affect the offsets of later edits.
+	sort.Slice(edits, func(i, j int) bool {
+		return edits[i].begin > edits[j].begin
+	})
+
+	for _, edit := range edits {
+		prefix := contents[:edit.begin]
+		suffix := contents[edit.end:]
+		contents = append(prefix, append(edit.replacement, suffix...)...)
+	}
+
+	return contents
+}
+
+// formatFuncDoc reformats the comment lines in commentList, and appends any
+// changes to the edit list.
+func formatFuncDoc(fileSet *token.FileSet, commentList []*ast.Comment, edits *edits) {
+	// Building the edit list to format a comment block is a two-step process.
+	// First, we iterate over each comment line looking for Swag attributes. In
+	// each one we find, we replace alignment whitespace with a tab character,
+	// then write the result into a tab writer.
+
+	linesToComments := make(map[int]int, len(commentList))
+
+	buffer := &bytes.Buffer{}
+	w := tabwriter.NewWriter(buffer, 1, 4, 1, '\t', 0)
+
+	for commentIndex, comment := range commentList {
+		text := comment.Text
+		if attr, body, found := swagComment(text); found {
+			formatted := "//\t" + attr
+			if body != "" {
+				formatted += "\t" + splitComment2(attr, body)
+			}
+			_, _ = fmt.Fprintln(w, formatted)
+			linesToComments[len(linesToComments)] = commentIndex
+		}
+	}
+
+	// Once we've loaded all of the comment lines to be aligned into the tab
+	// writer, flushing it causes the aligned text to be written out to the
+	// backing buffer.
+	_ = w.Flush()
+
+	// Now the second step: we iterate over the aligned comment lines that were
+	// written into the backing buffer, pair each one up to its original
+	// comment line, and use the combination to describe the edit that needs to
+	// be made to the original input.
+	formattedComments := bytes.Split(buffer.Bytes(), []byte("\n"))
+	for lineIndex, commentIndex := range linesToComments {
+		comment := commentList[commentIndex]
+		*edits = append(*edits, edit{
+			begin:       fileSet.Position(comment.Pos()).Offset,
+			end:         fileSet.Position(comment.End()).Offset,
+			replacement: formattedComments[lineIndex],
+		})
+	}
+}
+
+func splitComment2(attr, body string) string {
+	if specialTagForSplit[strings.ToLower(attr)] {
+		for i := 0; i < len(body); i++ {
+			if skipEnd, ok := skipChar[body[i]]; ok {
+				skipStart, n := body[i], 1
+				for i++; i < len(body); i++ {
+					if skipStart != skipEnd && body[i] == skipStart {
+						n++
+					} else if body[i] == skipEnd {
+						n--
+						if n == 0 {
+							break
+						}
+					}
+				}
+			} else if body[i] == ' ' || body[i] == '\t' {
+				j := i
+				for ; j < len(body) && (body[j] == ' ' || body[j] == '\t'); j++ {
+				}
+				body = replaceRange(body, i, j, "\t")
+			}
+		}
+	}
+	return body
+}
+
+func replaceRange(s string, start, end int, new string) string {
+	return s[:start] + new + s[end:]
+}
+
+var swagCommentLineExpression = regexp.MustCompile(`^\/\/\s+(@[\S.]+)\s*(.*)`)
+
+func swagComment(comment string) (string, string, bool) {
+	matches := swagCommentLineExpression.FindStringSubmatch(comment)
+	if matches == nil {
+		return "", "", false
+	}
+	return matches[1], matches[2], true
+}