Add optional formating for custom generator.

cmd/codegen/codegen.go

@@ -25,9 +25,10 @@ func main() {
 	generators := flag.String("generators", "all", "Comma separated list of generators to run. Options: all, model, convert, custom.")
 	// Convert flags
 	withTest := flag.Bool("convert.with-test", false, "Generate test functions for conversion functions. Default is true.")
-	// GQL flags
-	gqlOutFile := flag.String("custom.output-file", "", "Path of the generate gql file that is appened to the outputDir.")
-	gqlTemplateFile := flag.String("custom.template-file", "", "Path to the template file. Which is executed with codegen.TemplateData data.")
+	// Custom flags
+	customOutFile := flag.String("custom.output-file", "", "Path of the generate gql file that is appened to the outputDir.")
+	customTemplateFile := flag.String("custom.template-file", "", "Path to the template file. Which is executed with codegen.TemplateData data.")
+	customFormat := flag.Bool("custom.format", false, "Format the generated file with goimports.")
 	// Migration flags
 	migrationFileName := flag.String("migration.file-name", "", "Name of the migration file. Default is the model name.")
 
@@ -64,8 +65,9 @@ func main() {
 		PackageName: *packageName,
 		OutputDir:   *outputDir,
 		Custom: custom.Config{
-			OutputFile:   *gqlOutFile,
-			TemplateFile: *gqlTemplateFile,
+			OutputFile:   *customOutFile,
+			TemplateFile: *customTemplateFile,
+			Format:       *customFormat,
 		},
 		Convert: convert.Config{
 			WithTest: *withTest,

internal/generator/custom/custom.go

@@ -2,13 +2,15 @@
 package custom
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path"
 	"path/filepath"
 	"strings"
 	"text/template"
 
+	"github.com/DIMO-Network/model-garage/pkg/codegen"
 	"github.com/DIMO-Network/model-garage/pkg/schema"
 )
 
@@ -18,9 +20,10 @@ var customFileFormat = "%s.txt"
 type Config struct {
 	// OutputFile is the name of the model to generate the custom file.
 	OutputFile string
-
 	// TemplateFile is the path to the template file.
 	TemplateFile string
+	// Format controls whether the generated file is formatted with goimports.
+	Format bool
 }
 
 // Generate creates a new Custom file.
@@ -49,10 +52,19 @@ func Generate(tmplData *schema.TemplateData, outputDir string, cfg Config) error
 		}
 	}()
 
-	err = customFileTmpl.Execute(customFileOutputFile, &tmplData)
+	var outBuf bytes.Buffer
+	err = customFileTmpl.Execute(&outBuf, &tmplData)
 	if err != nil {
 		return fmt.Errorf("error executing Custom template: %w", err)
 	}
+	if cfg.Format {
+		err = codegen.FormatAndWriteToFile(outBuf.Bytes(), filePath)
+		if err != nil {
+			return fmt.Errorf("error writing file: %w", err)
+		}
+		return nil
+	}
+	customFileOutputFile.Write(outBuf.Bytes())
 
 	return nil
 }