generate.go

package main

import (
	"bytes"
	"fmt"
	"github.com/pkg/errors"
	"golang.org/x/tools/imports"
	"io"
	"log"
	"sort"
	"strings"
	"unicode"
)

func writeImports(w io.Writer, imports []string) {
	if imports == nil || len(imports) == 0 {
		return
	}

	fmt.Fprintln(w, "import (")

	for _, i := range imports {
		fmt.Fprintln(w, i)
	}

	fmt.Fprintln(w, ")")
}

func bool2int(b bool) int {
	if b {
		return 1
	}
	return 0
}

func buildInputParams(fields []genField) string {
	var fieldList []string

	// sort fields, required first
	sort.SliceStable(fields, func(i, j int) bool {
		return bool2int(fields[i].optional) < bool2int(fields[j].optional)
	})

	for _, f := range fields {
		if f.skip {
			continue
		}

		var optionalStr string
		if f.optional && !f.array {
			optionalStr = "*"
		}

		if f.array {
			if f.ptr {
				fieldList = append(fieldList, fmt.Sprintf("%s %s[]*%s", f.name, optionalStr, f.typ))
			} else {
				fieldList = append(fieldList, fmt.Sprintf("%s %s[]%s", f.name, optionalStr, f.typ))
			}
		} else {
			fieldList = append(fieldList, fmt.Sprintf("%s %s%s", f.name, optionalStr, f.typ))
		}
	}

	return strings.Join(fieldList, ",")
}

func buildBody(name string, fields []genField) string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("result := &%s {\n", name))

	// process required fields
	for _, f := range fields {
		if f.skip || f.optional {
			continue
		}

		// if field is required, not an array, and a pointer, then set to the address of the input param
		if f.ptr && !f.array {
			sb.WriteString(fmt.Sprintf("%s: &%s,\n", f.name, f.name))
		} else {
			sb.WriteString(fmt.Sprintf("%s: %s,\n", f.name, f.name))
		}
	}

	sb.WriteString("}\n")

	// process optional fields
	for _, f := range fields {
		if f.skip || !f.optional {
			continue
		}

		if f.ptr {
			sb.WriteString(fmt.Sprintf("if %s != nil {\nresult.%s = %s\n}\n", f.name, f.name, f.name))
		} else {
			sb.WriteString(fmt.Sprintf("if %s != nil {\nresult.%s = *%s\n}\n", f.name, f.name, f.name))
		}
	}

	sb.WriteString("return result\n")

	return sb.String()
}

func formatStructName(in string) string {
	return "New" + string(unicode.ToUpper(rune(in[0]))) + in[1:]
}

func writeStruct(w io.Writer, s genStruct) {
	fmFuncName := formatStructName(s.name)
	comment := fmt.Sprintf("// %s generated factory method for %s", fmFuncName, s.name)

	// struct comment
	fmt.Fprintln(w, comment)

	// struct method signature
	fmt.Fprintf(w, "func %s(%s) *%s{\n", fmFuncName, buildInputParams(s.fields), s.name)

	// build struct body
	fmt.Fprintf(w, buildBody(s.name, s.fields))
	fmt.Fprintln(w, "}")
}

func writePackageFile(w io.Writer, pkg string, pkgImports []string, structs []genStruct) {
	var buf bytes.Buffer
	var err error

	// write the header comment
	fmt.Fprintln(&buf, `// Code generated by "fmgen". DO NOT EDIT.`)

	// write the package
	fmt.Fprintln(&buf, fmt.Sprintf("package %s", pkg))

	// write the imports
	writeImports(&buf, pkgImports)

	// write factory methods for each struct
	for _, s := range structs {
		if !s.Skip() {
			writeStruct(&buf, s)
		}
	}

	output := fmt.Sprintf("%s/%s", pkg, generatedFileName)
	log.Printf("generating factory method file for package [%s]", pkg)

	opts := &imports.Options{
		Fragment:   false,
		AllErrors:  true,
		Comments:   true,
		TabIndent:  false,
		TabWidth:   4,
		FormatOnly: false,
	}

	// process imports twice due to issues noticed with imports not being removed
	processed := buf.Bytes()
	for i := 0; i < 2; i++ {
		processed, err = imports.Process(output, processed, opts)
		if err != nil {
			log.Panicf("unable to process generated file to remove unused imports - %v", errors.WithStack(err))
		}
	}

	_, err = io.Copy(w, bytes.NewReader(processed))
	if err != nil {
		log.Panicf("unable to copy processed results to writer - %v", errors.WithStack(err))
	}
}