niki/vendor/github.com/swaggo/swag/formatter.go

183 lines
4.9 KiB
Go
Raw Permalink Normal View History

2024-05-14 13:07:09 +00:00
package swag
import (
"bytes"
"fmt"
"go/ast"
goparser "go/parser"
"go/token"
"log"
"os"
"regexp"
"sort"
"strings"
"text/tabwriter"
)
// 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)
}
return edits.apply(contents), 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
}