formatter.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. package swag
  2. import (
  3. "bytes"
  4. "fmt"
  5. "go/ast"
  6. goparser "go/parser"
  7. "go/token"
  8. "log"
  9. "os"
  10. "regexp"
  11. "sort"
  12. "strings"
  13. "text/tabwriter"
  14. )
  15. // Check of @Param @Success @Failure @Response @Header
  16. var specialTagForSplit = map[string]bool{
  17. paramAttr: true,
  18. successAttr: true,
  19. failureAttr: true,
  20. responseAttr: true,
  21. headerAttr: true,
  22. }
  23. var skipChar = map[byte]byte{
  24. '"': '"',
  25. '(': ')',
  26. '{': '}',
  27. '[': ']',
  28. }
  29. // Formatter implements a formatter for Go source files.
  30. type Formatter struct {
  31. // debugging output goes here
  32. debug Debugger
  33. }
  34. // NewFormatter create a new formatter instance.
  35. func NewFormatter() *Formatter {
  36. formatter := &Formatter{
  37. debug: log.New(os.Stdout, "", log.LstdFlags),
  38. }
  39. return formatter
  40. }
  41. // Format formats swag comments in contents. It uses fileName to report errors
  42. // that happen during parsing of contents.
  43. func (f *Formatter) Format(fileName string, contents []byte) ([]byte, error) {
  44. fileSet := token.NewFileSet()
  45. ast, err := goparser.ParseFile(fileSet, fileName, contents, goparser.ParseComments)
  46. if err != nil {
  47. return nil, err
  48. }
  49. // Formatting changes are described as an edit list of byte range
  50. // replacements. We make these content-level edits directly rather than
  51. // changing the AST nodes and writing those out (via [go/printer] or
  52. // [go/format]) so that we only change the formatting of Swag attribute
  53. // comments. This won't touch the formatting of any other comments, or of
  54. // functions, etc.
  55. maxEdits := 0
  56. for _, comment := range ast.Comments {
  57. maxEdits += len(comment.List)
  58. }
  59. edits := make(edits, 0, maxEdits)
  60. for _, comment := range ast.Comments {
  61. formatFuncDoc(fileSet, comment.List, &edits)
  62. }
  63. return edits.apply(contents), nil
  64. }
  65. type edit struct {
  66. begin int
  67. end int
  68. replacement []byte
  69. }
  70. type edits []edit
  71. func (edits edits) apply(contents []byte) []byte {
  72. // Apply the edits with the highest offset first, so that earlier edits
  73. // don't affect the offsets of later edits.
  74. sort.Slice(edits, func(i, j int) bool {
  75. return edits[i].begin > edits[j].begin
  76. })
  77. for _, edit := range edits {
  78. prefix := contents[:edit.begin]
  79. suffix := contents[edit.end:]
  80. contents = append(prefix, append(edit.replacement, suffix...)...)
  81. }
  82. return contents
  83. }
  84. // formatFuncDoc reformats the comment lines in commentList, and appends any
  85. // changes to the edit list.
  86. func formatFuncDoc(fileSet *token.FileSet, commentList []*ast.Comment, edits *edits) {
  87. // Building the edit list to format a comment block is a two-step process.
  88. // First, we iterate over each comment line looking for Swag attributes. In
  89. // each one we find, we replace alignment whitespace with a tab character,
  90. // then write the result into a tab writer.
  91. linesToComments := make(map[int]int, len(commentList))
  92. buffer := &bytes.Buffer{}
  93. w := tabwriter.NewWriter(buffer, 0, 0, 1, ' ', 0)
  94. for commentIndex, comment := range commentList {
  95. text := comment.Text
  96. if attr, body, found := swagComment(text); found {
  97. formatted := "// " + attr
  98. if body != "" {
  99. formatted += "\t" + splitComment2(attr, body)
  100. }
  101. _, _ = fmt.Fprintln(w, formatted)
  102. linesToComments[len(linesToComments)] = commentIndex
  103. }
  104. }
  105. // Once we've loaded all of the comment lines to be aligned into the tab
  106. // writer, flushing it causes the aligned text to be written out to the
  107. // backing buffer.
  108. _ = w.Flush()
  109. // Now the second step: we iterate over the aligned comment lines that were
  110. // written into the backing buffer, pair each one up to its original
  111. // comment line, and use the combination to describe the edit that needs to
  112. // be made to the original input.
  113. formattedComments := bytes.Split(buffer.Bytes(), []byte("\n"))
  114. for lineIndex, commentIndex := range linesToComments {
  115. comment := commentList[commentIndex]
  116. *edits = append(*edits, edit{
  117. begin: fileSet.Position(comment.Pos()).Offset,
  118. end: fileSet.Position(comment.End()).Offset,
  119. replacement: formattedComments[lineIndex],
  120. })
  121. }
  122. }
  123. func splitComment2(attr, body string) string {
  124. if specialTagForSplit[strings.ToLower(attr)] {
  125. for i := 0; i < len(body); i++ {
  126. if skipEnd, ok := skipChar[body[i]]; ok {
  127. if skipLen := strings.IndexByte(body[i+1:], skipEnd); skipLen > 0 {
  128. i += skipLen
  129. }
  130. } else if body[i] == ' ' {
  131. j := i
  132. for ; j < len(body) && body[j] == ' '; j++ {
  133. }
  134. body = replaceRange(body, i, j, "\t")
  135. }
  136. }
  137. }
  138. return body
  139. }
  140. func replaceRange(s string, start, end int, new string) string {
  141. return s[:start] + new + s[end:]
  142. }
  143. var swagCommentLineExpression = regexp.MustCompile(`^\/\/\s+(@[\S.]+)\s*(.*)`)
  144. func swagComment(comment string) (string, string, bool) {
  145. matches := swagCommentLineExpression.FindStringSubmatch(comment)
  146. if matches == nil {
  147. return "", "", false
  148. }
  149. return matches[1], matches[2], true
  150. }