Refactor output

output.go

@@ -19,55 +19,87 @@ var testHarnessFileOpen = func(filename string) (io.WriteCloser, error) {
 
 // generateOutput builds the file output and sends it to outHandler for saving
 func generateOutput(state *State, data *templateData) error {
-	if len(state.Templates) == 0 {
+	return executeTemplates(executeTemplateData{
-		return errors.New("No template files located for generation")
+		state:                state,
-	}
+		data:                 data,
-	var out [][]byte
+		templates:            state.Templates,
-	var imps imports
+		importSet:            defaultTemplateImports,
-
+		combineImportsOnType: true,
-	imps.standard = defaultTemplateImports.standard
+		fileSuffix:           ".go",
-	imps.thirdParty = defaultTemplateImports.thirdParty
+	})
-
-	for _, template := range state.Templates {
-		imps = combineTypeImports(imps, importsBasedOnType, data.Table.Columns)
-		resp, err := generateTemplate(template, data)
-		if err != nil {
-			return fmt.Errorf("Error generating template %s: %s", template.Name(), err)
-		}
-		out = append(out, resp)
-	}
-
-	fName := data.Table.Name + ".go"
-	err := outHandler(state.Config.OutFolder, fName, state.Config.PkgName, imps, out)
-	if err != nil {
-		return err
-	}
-
-	return nil
 }
 
 // generateTestOutput builds the test file output and sends it to outHandler for saving
 func generateTestOutput(state *State, data *templateData) error {
-	if len(state.TestTemplates) == 0 {
+	return executeTemplates(executeTemplateData{
-		return errors.New("No template files located for generation")
+		state:                state,
-	}
+		data:                 data,
+		templates:            state.TestTemplates,
+		importSet:            defaultTestTemplateImports,
+		combineImportsOnType: false,
+		fileSuffix:           "_test.go",
+	})
+}
 
+// generateSingletonOutput processes the templates that should only be run
+// one time.
+func generateSingletonOutput(state *State, data *templateData) error {
+	return executeSingletonTemplates(executeTemplateData{
+		state:          state,
+		data:           data,
+		templates:      state.SingletonTemplates,
+		importNamedSet: defaultSingletonTemplateImports,
+		fileSuffix:     ".go",
+	})
+}
+
+// generateSingletonTestOutput processes the templates that should only be run
+// one time.
+func generateSingletonTestOutput(state *State, data *templateData) error {
+	return executeSingletonTemplates(executeTemplateData{
+		state:          state,
+		data:           data,
+		templates:      state.SingletonTestTemplates,
+		importNamedSet: defaultSingletonTestTemplateImports,
+		fileSuffix:     "_test.go",
+	})
+}
+
+type executeTemplateData struct {
+	state *State
+	data  *templateData
+
+	templates templateList
+
+	importSet      imports
+	importNamedSet map[string]imports
+
+	combineImportsOnType bool
+
+	fileSuffix string
+}
+
+func executeTemplates(e executeTemplateData) error {
 	var out [][]byte
 	var imps imports
 
-	imps.standard = defaultTestTemplateImports.standard
+	imps.standard = e.importSet.standard
-	imps.thirdParty = defaultTestTemplateImports.thirdParty
+	imps.thirdParty = e.importSet.thirdParty
 
-	for _, template := range state.TestTemplates {
+	for _, template := range e.templates {
-		resp, err := generateTemplate(template, data)
+		if e.combineImportsOnType {
+			imps = combineTypeImports(imps, importsBasedOnType, e.data.Table.Columns)
+		}
+
+		resp, err := executeTemplate(template, e.data)
 		if err != nil {
-			return fmt.Errorf("Error generating test template %s: %s", template.Name(), err)
+			return fmt.Errorf("Error generating template %s: %s", template.Name(), err)
 		}
 		out = append(out, resp)
 	}
 
-	fName := data.Table.Name + "_test.go"
+	fName := e.data.Table.Name + e.fileSuffix
-	err := outHandler(state.Config.OutFolder, fName, state.Config.PkgName, imps, out)
+	err := outHandler(e.state.Config.OutFolder, fName, e.state.Config.PkgName, imps, out)
 	if err != nil {
 		return err
 	}
@@ -75,22 +107,9 @@ func generateTestOutput(state *State, data *templateData) error {
 	return nil
 }
 
-// generateSingletonOutput processes the templates that should only be run
+func executeSingletonTemplates(e executeTemplateData) error {
-// one time.
+	for _, template := range e.templates {
-func generateSingletonOutput(state *State) error {
+		resp, err := executeTemplate(template, e.data)
-	if state.SingletonTemplates == nil {
-		return errors.New("No singleton templates located for generation")
-	}
-
-	templateData := &templateData{
-		PkgName:    state.Config.PkgName,
-		DriverName: state.Config.DriverName,
-	}
-
-	for _, template := range state.SingletonTemplates {
-		var imps imports
-
-		resp, err := generateTemplate(template, templateData)
 		if err != nil {
 			return fmt.Errorf("Error generating template %s: %s", template.Name(), err)
 		}
@@ -99,12 +118,18 @@ func generateSingletonOutput(state *State) error {
 		ext := filepath.Ext(fName)
 		fName = fName[0 : len(fName)-len(ext)]
 
-		imps.standard = defaultSingletonTemplateImports[fName].standard
+		imps := imports{
-		imps.thirdParty = defaultSingletonTemplateImports[fName].thirdParty
+			standard:   e.importNamedSet[fName].standard,
+			thirdParty: e.importNamedSet[fName].thirdParty,
+		}
 
-		fName = fName + ".go"
+		err = outHandler(
-
+			e.state.Config.OutFolder,
-		err = outHandler(state.Config.OutFolder, fName, state.Config.PkgName, imps, [][]byte{resp})
+			fName+e.fileSuffix,
+			e.state.Config.PkgName,
+			imps,
+			[][]byte{resp},
+		)
 		if err != nil {
 			return err
 		}
@@ -113,45 +138,7 @@ func generateSingletonOutput(state *State) error {
 	return nil
 }
 
-// generateSingletonTestOutput processes the templates that should only be run
+func generateTestMainOutput(state *State, data *templateData) error {
-// one time.
-func generateSingletonTestOutput(state *State) error {
-	if state.SingletonTestTemplates == nil {
-		return errors.New("No singleton test templates located for generation")
-	}
-
-	templateData := &templateData{
-		PkgName:    state.Config.PkgName,
-		DriverName: state.Config.DriverName,
-	}
-
-	for _, template := range state.SingletonTestTemplates {
-		var imps imports
-
-		resp, err := generateTemplate(template, templateData)
-		if err != nil {
-			return fmt.Errorf("Error generating test template %s: %s", template.Name(), err)
-		}
-
-		fName := template.Name()
-		ext := filepath.Ext(fName)
-		fName = fName[0 : len(fName)-len(ext)]
-
-		imps.standard = defaultSingletonTestTemplateImports[fName].standard
-		imps.thirdParty = defaultSingletonTestTemplateImports[fName].thirdParty
-
-		fName = fName + "_test.go"
-
-		err = outHandler(state.Config.OutFolder, fName, state.Config.PkgName, imps, [][]byte{resp})
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func generateTestMainOutput(state *State) error {
 	if state.TestMainTemplate == nil {
 		return errors.New("No TestMain template located for generation")
 	}
@@ -162,13 +149,7 @@ func generateTestMainOutput(state *State) error {
 	imps.standard = defaultTestMainImports[state.Config.DriverName].standard
 	imps.thirdParty = defaultTestMainImports[state.Config.DriverName].thirdParty
 
-	templateData := &templateData{
+	resp, err := executeTemplate(state.TestMainTemplate, data)
-		Tables:     state.Tables,
-		PkgName:    state.Config.PkgName,
-		DriverName: state.Config.DriverName,
-	}
-
-	resp, err := generateTemplate(state.TestMainTemplate, templateData)
 	if err != nil {
 		return err
 	}
@@ -214,8 +195,9 @@ func outHandler(outFolder string, fileName string, pkgName string, imps imports,
 	return nil
 }
 
-// generateTemplate takes a template and returns the output of the template execution.
+// executeTemplate takes a template and returns the output of the template
-func generateTemplate(t *template.Template, data *templateData) ([]byte, error) {
+// execution.
+func executeTemplate(t *template.Template, data *templateData) ([]byte, error) {
 	var buf bytes.Buffer
 	if err := t.Execute(&buf, data); err != nil {
 		return nil, err

sqlboiler.go

@@ -70,16 +70,22 @@ func New(config *Config) (*State, error) {
 // Run executes the sqlboiler templates and outputs them to files based on the
 // state given.
 func (s *State) Run(includeTests bool) error {
-	if err := generateSingletonOutput(s); err != nil {
+	singletonData := &templateData{
+		Tables:     s.Tables,
+		DriverName: s.Config.DriverName,
+		PkgName:    s.Config.PkgName,
+	}
+
+	if err := generateSingletonOutput(s, singletonData); err != nil {
 		return fmt.Errorf("Unable to generate singleton template output: %s", err)
 	}
 
 	if includeTests {
-		if err := generateTestMainOutput(s); err != nil {
+		if err := generateTestMainOutput(s, singletonData); err != nil {
 			return fmt.Errorf("Unable to generate TestMain output: %s", err)
 		}
 
-		if err := generateSingletonTestOutput(s); err != nil {
+		if err := generateSingletonTestOutput(s, singletonData); err != nil {
 			return fmt.Errorf("Unable to generate singleton test template output: %s", err)
 		}
 	}