fix: update 2 times with dstFile

haunt98 · Jan 17, 2023 · 6bbfd7b · 6bbfd7b
1 parent 031a349
commit 6bbfd7b
Showing 1 changed file with 21 additions and 1 deletion.
diff --git a/internal/imports/formatter.go b/internal/imports/formatter.go
@@ -234,7 +234,7 @@ func (ft *Formatter) formatImports(
 	if err != nil {
 		return nil, fmt.Errorf("decorator: failed to parse file [%s]: %w", path, err)
 	}
-	if len(dstFile.Imports) == 0 {
+	if len(dstFile.Imports) == 0 || len(dstFile.Decls) == 0 {
 		return nil, ErrEmptyImport
 	}
 	ft.logDSTImportSpecs("formatImports: dstImportSpecs", dstFile.Imports)
@@ -250,8 +250,28 @@ func (ft *Formatter) formatImports(
 	}
 	ft.logDSTImportSpecs("formatImports: formattedDSTImportSpecs: ", formattedDSTImportSpecs)
 
+	// First update
 	dstFile.Imports = formattedDSTImportSpecs
 
+	genSpecs := dstFile.Decls[0].(*dst.GenDecl).Specs
+	formattedGenSpecs := make([]dst.Spec, 0, len(genSpecs))
+
+	// Append all imports first
+	for _, importSpec := range formattedDSTImportSpecs {
+		formattedGenSpecs = append(formattedGenSpecs, importSpec)
+	}
+
+	// Append all non imports later
+	for _, genSpec := range genSpecs {
+		if _, ok := genSpec.(*dst.ImportSpec); !ok {
+			formattedGenSpecs = append(formattedGenSpecs, genSpec)
+			continue
+		}
+	}
+
+	// Second update
+	dstFile.Decls[0].(*dst.GenDecl).Specs = formattedGenSpecs
+
 	var buf bytes.Buffer
 	if err := decorator.Fprint(&buf, dstFile); err != nil {
 		return nil, fmt.Errorf("decorator: failed to fprint [%s]: %w", path, err)