formatting for struct (#52)

reviser/reviser.go

@@ -122,25 +122,33 @@ func formatDecls(f *ast.File, options Options) {
 	}
 
 	for _, decl := range f.Decls {
-		if dd, ok := decl.(*ast.FuncDecl); ok {
-			var formattedComments []*ast.Comment
-			if dd.Doc != nil {
-				formattedComments = make([]*ast.Comment, len(dd.Doc.List))
-			}
+		switch dd := decl.(type) {
+		case *ast.GenDecl:
+			dd.Doc = fixCommentGroup(dd.Doc)
+		case *ast.FuncDecl:
+			dd.Doc = fixCommentGroup(dd.Doc)
+		}
+	}
+}
 
-			formattedDoc := &ast.CommentGroup{
-				List: formattedComments,
-			}
+func fixCommentGroup(commentGroup *ast.CommentGroup) *ast.CommentGroup {
+	if commentGroup == nil {
+		formattedDoc := &ast.CommentGroup{
+			List: []*ast.Comment{},
+		}
 
-			if dd.Doc != nil {
-				for i, comment := range dd.Doc.List {
-					formattedDoc.List[i] = comment
-				}
-			}
+		return formattedDoc
+	}
 
-			dd.Doc = formattedDoc
-		}
+	formattedDoc := &ast.CommentGroup{
+		List: make([]*ast.Comment, len(commentGroup.List)),
 	}
+
+	for i, comment := range commentGroup.List {
+		formattedDoc.List[i] = comment
+	}
+
+	return formattedDoc
 }
 
 func groupImports(

reviser/reviser_test.go

@@ -954,12 +954,26 @@ func TestExecute_WithFormat(t *testing.T) {
 				projectName: "github.com/incu6us/goimports-reviser",
 				filePath:    "./testdata/example.go",
 				fileContent: `package testdata
+type SomeStruct struct{}
+type SomeStruct1 struct{}
+// SomeStruct2 comments
+type SomeStruct2 struct{}
+func (s *SomeStruct2) test() {}
 func test(){}
 func test1(){}
 `,
 			},
 			want: `package testdata
 
+type SomeStruct struct{}
+
+type SomeStruct1 struct{}
+
+// SomeStruct2 comments
+type SomeStruct2 struct{}
+
+func (s *SomeStruct2) test() {}
+
 func test() {}
 
 func test1() {}