From 9df66a9b574a280d3979e25aee476f1793f06c5a Mon Sep 17 00:00:00 2001
From: Kevin Joiner <10265309+KevinJoiner@users.noreply.github.com>
Date: Thu, 23 May 2024 14:33:28 -0400
Subject: [PATCH] Add optional formating for custom generator.

---
 cmd/codegen/codegen.go              | 12 +++++++-----
 internal/generator/custom/custom.go | 16 ++++++++++++++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/cmd/codegen/codegen.go b/cmd/codegen/codegen.go
index 0845e19..7eec7a0 100644
--- a/cmd/codegen/codegen.go
+++ b/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,
diff --git a/internal/generator/custom/custom.go b/internal/generator/custom/custom.go
index a0e3637..b0b6170 100644
--- a/internal/generator/custom/custom.go
+++ b/internal/generator/custom/custom.go
@@ -2,6 +2,7 @@
 package custom
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path"
@@ -9,6 +10,7 @@ import (
 	"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
 }