forked from ebhomengo/niki
183 lines
4.9 KiB
Go
183 lines
4.9 KiB
Go
|
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
|
||
|
}
|