output.go

package codegen

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"

	"github.com/pkg/errors"
)

func Output(ctx *GenContext, filePath string) error {
	if ctx.PackageName == "" {
		return errors.New("missing package name")
	}

	var unformatted bytes.Buffer
	fmt.Fprint(&unformatted, "// Code generated by go-codegen; DO NOT EDIT.\n\n")
	fmt.Fprintf(&unformatted, "package %s\n", ctx.PackageName)
	outputImports(ctx.Imports(), &unformatted)
	for _, g := range ctx.Generated() {
		fmt.Fprintln(&unformatted, g)
	}

	// Parsing and printing ensures formatting.
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, filePath, unformatted.Bytes(), parser.ParseComments)
	if err != nil {
		fmt.Println(unformatted.String())
		return errors.Wrap(err, "parsing generated code")
	}

	var formatted bytes.Buffer
	printer.Fprint(&formatted, fset, file)
	return ioutil.WriteFile(filePath, formatted.Bytes(), 0644)
}

func outputImports(imports []string, w io.Writer) {
	if len(imports) == 0 {
		return
	}
	fmt.Fprintln(w, "import (")
	for _, i := range imports {
		fmt.Fprintf(w, "\t\"%s\"\n", i)
	}
	fmt.Fprintln(w, ")")
}