chore: Move file format handling code to format/format.go

- Handle file open/read/update in format/format.go - Only handle formatting file contents in formatter.go - Write tests against formatter public API only - Clean up formatter code
swaggo · Apr 28, 2022 · 9ceda39 · 9ceda39
1 parent 0077fac
commit 9ceda39
Show file tree

Hide file tree

Showing 4 changed files with 294 additions and 376 deletions.
diff --git a/format/format.go b/format/format.go
@@ -2,6 +2,7 @@ package format
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
@@ -60,24 +61,52 @@ func (f *Format) Build(config *Config) error {
 }
 
 func (f *Format) visit(path string, fileInfo os.FileInfo, err error) error {
-	if fileInfo.IsDir() {
-		return f.skipDir(path, fileInfo)
+	if fileInfo.IsDir() && f.excludeDir(path) {
+		return filepath.SkipDir
 	}
-	if f.exclude[path] ||
-		strings.HasSuffix(strings.ToLower(path), "_test.go") ||
-		filepath.Ext(path) != ".go" {
+	if f.excludeFile(path) {
 		return nil
 	}
-	if err := f.formatter.Format(path); err != nil {
+	if err := f.format(path); err != nil {
 		return fmt.Errorf("fmt: %w", err)
 	}
 	return nil
 }
 
-func (f *Format) skipDir(path string, info os.FileInfo) error {
-	if f.exclude[path] ||
-		len(info.Name()) > 1 && info.Name()[0] == '.' { // exclude hidden folders
-		return filepath.SkipDir
+func (f *Format) excludeDir(path string) bool {
+	return f.exclude[path] ||
+		filepath.Base(path)[0] == '.' && len(filepath.Base(path)) > 1 // exclude hidden folders
+}
+
+func (f *Format) excludeFile(path string) bool {
+	return f.exclude[path] ||
+		strings.HasSuffix(strings.ToLower(path), "_test.go") ||
+		filepath.Ext(path) != ".go"
+}
+
+func (f *Format) format(path string) error {
+	contents, err := ioutil.ReadFile(path)
+	if err != nil {
+		return err
 	}
-	return nil
+	formatted, err := f.formatter.Format(path, contents)
+	if err != nil {
+		return err
+	}
+	return write(path, formatted)
+}
+
+func write(path string, contents []byte) error {
+	f, err := ioutil.TempFile(filepath.Split(path))
+	if err != nil {
+		return err
+	}
+	defer os.Remove(f.Name())
+	if _, err := f.Write(contents); err != nil {
+		return err
+	}
+	if err := f.Close(); err != nil {
+		return err
+	}
+	return os.Rename(f.Name(), path)
 }
diff --git a/format/format_test.go b/format/format_test.go
@@ -49,6 +49,19 @@ func TestFormat_ParseError(t *testing.T) {
 	assert.Error(t, New().Build(&Config{SearchDir: fx.basedir}))
 }
 
+func TestFormat_ReadError(t *testing.T) {
+	fx := setup(t)
+	os.Chmod(filepath.Join(fx.basedir, "main.go"), 0)
+	assert.Error(t, New().Build(&Config{SearchDir: fx.basedir}))
+}
+
+func TestFormat_WriteError(t *testing.T) {
+	fx := setup(t)
+	os.Chmod(fx.basedir, 0555)
+	assert.Error(t, New().Build(&Config{SearchDir: fx.basedir}))
+	os.Chmod(fx.basedir, 0755)
+}
+
 func TestFormat_InvalidSearchDir(t *testing.T) {
 	formatter := New()
 	assert.Error(t, formatter.Build(&Config{SearchDir: "no_such_dir"}))

diff --git a/formatter.go b/formatter.go
@@ -8,17 +8,31 @@ import (
 	goparser "go/parser"
 	"go/token"
 	"io"
-	"io/ioutil"
 	"log"
 	"os"
-	"path/filepath"
 	"regexp"
 	"strings"
 	"text/tabwriter"
 )
 
 const splitTag = "&*"
 
+// 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
@@ -33,199 +47,83 @@ func NewFormatter() *Formatter {
 	return formatter
 }
 
-// Format swag comments in given file.
-func (f *Formatter) Format(filepath string) error {
+func (f *Formatter) Format(fileName string, contents []byte) ([]byte, error) {
 	fileSet := token.NewFileSet()
-	astFile, err := goparser.ParseFile(fileSet, filepath, nil, goparser.ParseComments)
+	ast, err := goparser.ParseFile(fileSet, fileName, contents, goparser.ParseComments)
 	if err != nil {
-		return err
+		return nil, err
 	}
+	formattedComments := bytes.Buffer{}
+	oldComments := map[string]string{}
 
-	var (
-		formatedComments = bytes.Buffer{}
-		// CommentCache
-		oldCommentsMap = make(map[string]string)
-	)
-
-	if astFile.Comments != nil {
-		for _, comment := range astFile.Comments {
-			formatFuncDoc(comment.List, &formatedComments, oldCommentsMap)
+	if ast.Comments != nil {
+		for _, comment := range ast.Comments {
+			formatFuncDoc(comment.List, &formattedComments, oldComments)
 		}
 	}
-
-	return writeFormattedComments(filepath, formatedComments, oldCommentsMap)
+	return formatComments(fileName, contents, formattedComments.Bytes(), oldComments), nil
 }
 
-func writeFormattedComments(filepath string, formatedComments bytes.Buffer, oldCommentsMap map[string]string) error {
-	// Replace the file
-	// Read the file
-	srcBytes, err := ioutil.ReadFile(filepath)
-	if err != nil {
-		return fmt.Errorf("cannot open file, err: %w path : %s ", err, filepath)
-	}
-
-	replaceSrc, newComments := string(srcBytes), strings.Split(formatedComments.String(), "\n")
-
-	for _, e := range newComments {
-		commentSplit := strings.Split(e, splitTag)
-		if len(commentSplit) == 2 {
-			commentHash, commentContent := commentSplit[0], commentSplit[1]
-
-			if !isBlankComment(commentContent) {
-				replaceSrc = strings.Replace(replaceSrc, oldCommentsMap[commentHash], commentContent, 1)
-			}
+func formatComments(fileName string, contents []byte, formattedComments []byte, oldComments map[string]string) []byte {
+	for _, comment := range bytes.Split(formattedComments, []byte("\n")) {
+		hash, content, found := bytes.Cut(comment, []byte(splitTag))
+		if found {
+			contents = bytes.Replace(contents, []byte(oldComments[string(hash)]), content, 1)
 		}
 	}
-	return writeBack(filepath, []byte(replaceSrc))
+	return contents
 }
 
 func formatFuncDoc(commentList []*ast.Comment, formattedComments io.Writer, oldCommentsMap map[string]string) {
-	tabWriter := tabwriter.NewWriter(formattedComments, 0, 0, 2, ' ', 0)
+	w := tabwriter.NewWriter(formattedComments, 0, 0, 2, ' ', 0)
 
 	for _, comment := range commentList {
-		commentLine := comment.Text
-		if isSwagComment(commentLine) || isBlankComment(commentLine) {
-			cmd5 := fmt.Sprintf("%x", md5.Sum([]byte(commentLine)))
-
-			// Find the separator and replace to \t
-			c := separatorFinder(commentLine, '\t')
-			oldCommentsMap[cmd5] = commentLine
-
+		text := comment.Text
+		if attr, body, found := swagComment(text); found {
+			cmd5 := fmt.Sprintf("%x", md5.Sum([]byte(text)))
+			oldCommentsMap[cmd5] = text
+
+			formatted := "// " + attr
+			if body != "" {
+				formatted += "\t" + splitComment2(attr, body)
+			}
 			// md5 + splitTag + srcCommentLine
 			// eg. xxx&*@Description get struct array
-			_, _ = fmt.Fprintln(tabWriter, cmd5+splitTag+c)
-		}
-	}
-	// format by tabWriter
-	_ = tabWriter.Flush()
-}
-
-func separatorFinder(comment string, replacer byte) string {
-	commentBytes, commentLine := []byte(comment), strings.TrimSpace(strings.TrimLeft(comment, "/"))
-
-	if len(commentLine) == 0 {
-		return ""
-	}
-
-	attribute := strings.Fields(commentLine)[0]
-	attrLen := strings.Index(comment, attribute) + len(attribute)
-	attribute = strings.ToLower(attribute)
-
-	var (
-		length = attrLen
-
-		// Check of @Param @Success @Failure @Response @Header.
-		specialTagForSplit = map[string]byte{
-			paramAttr:    1,
-			successAttr:  1,
-			failureAttr:  1,
-			responseAttr: 1,
-			headerAttr:   1,
+			_, _ = fmt.Fprintln(w, cmd5+splitTag+formatted)
 		}
-	)
-
-	_, ok := specialTagForSplit[attribute]
-	if ok {
-		return splitSpecialTags(commentBytes, length, replacer)
-	}
-
-	for length < len(commentBytes) && commentBytes[length] == ' ' {
-		length++
-	}
-
-	if length >= len(commentBytes) {
-		return comment
 	}
-
-	commentBytes = replaceRange(commentBytes, attrLen, length, replacer)
-
-	return string(commentBytes)
+	// format by tabwriter
+	_ = w.Flush()
 }
 
-func splitSpecialTags(commentBytes []byte, length int, rp byte) string {
-	var (
-		skipFlag bool
-		skipChar = map[byte]byte{
-			'"': 1,
-			'(': 1,
-			'{': 1,
-			'[': 1,
-		}
-
-		skipCharEnd = map[byte]byte{
-			'"': 1,
-			')': 1,
-			'}': 1,
-			']': 1,
-		}
-	)
-
-	for ; length < len(commentBytes); length++ {
-		if !skipFlag && commentBytes[length] == ' ' {
-			j := length
-			for j < len(commentBytes) && commentBytes[j] == ' ' {
-				j++
+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 {
+				if skipLen := strings.IndexByte(body[i+1:], skipEnd); skipLen > 0 {
+					i += skipLen
+				}
+			} else if body[i] == ' ' {
+				j := i
+				for ; j < len(body) && body[j] == ' '; j++ {
+				}
+				body = replaceRange(body, i, j, "\t")
 			}
-
-			commentBytes = replaceRange(commentBytes, length, j, rp)
-		}
-
-		_, found := skipChar[commentBytes[length]]
-		if found && !skipFlag {
-			skipFlag = true
-
-			continue
 		}
-
-		_, found = skipCharEnd[commentBytes[length]]
-		if found && skipFlag {
-			skipFlag = false
-		}
-	}
-
-	return string(commentBytes)
-}
-
-func replaceRange(s []byte, start, end int, new byte) []byte {
-	if start > end || end < 1 {
-		return s
 	}
-
-	if end > len(s) {
-		end = len(s)
-	}
-
-	s = append(s[:start], s[end-1:]...)
-
-	s[start] = new
-
-	return s
+	return body
 }
 
-var swagCommentExpression = regexp.MustCompile("@[A-z]+")
-
-func isSwagComment(comment string) bool {
-	return swagCommentExpression.MatchString(strings.ToLower(comment))
+func replaceRange(s string, start, end int, new string) string {
+	return s[:start] + new + s[end:]
 }
 
-func isBlankComment(comment string) bool {
-	return len(strings.TrimSpace(comment)) == 0
-}
+var swagCommentLineExpression = regexp.MustCompile(`^\/\/\s+(@[\S.]+)\s*(.*)`)
 
-func writeBack(filename string, src []byte) error {
-	f, err := ioutil.TempFile(filepath.Dir(filename), filepath.Base(filename))
-	if err != nil {
-		return err
-	}
-	defer os.Remove(f.Name())
-	if _, err := f.Write(src); err != nil {
-		return err
-	}
-	if err := f.Close(); err != nil {
-		return err
-	}
-	if err := os.Rename(f.Name(), filename); err != nil {
-		return err
+func swagComment(comment string) (string, string, bool) {
+	matches := swagCommentLineExpression.FindStringSubmatch(comment)
+	if matches == nil {
+		return "", "", false
 	}
-	return nil
+	return matches[1], matches[2], true
 }