main.go

package main

import (
	"fmt"
	"go/types"
	"os"
	"strings"
	"unicode"
	"unicode/utf8"

	. "github.com/dave/jennifer/jen"
	"golang.org/x/tools/go/packages"
)

func main() {
	// 1. Handle arguments to command
	if len(os.Args) != 2 {
		failErr(fmt.Errorf("expected exactly one argument: <source type>"))
	}
	sourceType := os.Args[1]
	sourceTypePackage, sourceTypeName := splitSourceType(sourceType)

	// 2. Inspect package and use type checker to infer imported types
	pkg := loadPackage(sourceTypePackage)
	fmt.Println("original package path: " + sourceTypePackage)

	// 3. Lookup the given source type name in the package declarations
	obj := pkg.Types.Scope().Lookup(sourceTypeName)
	if obj == nil {
		failErr(fmt.Errorf("%s not found in declared types of %s",
			sourceTypeName, pkg))
	}

	// 4. We check if it is a declared type
	if _, ok := obj.(*types.TypeName); !ok {
		failErr(fmt.Errorf("%v is not a named type", obj))
	}
	// 5. We expect the underlying type to be a struct
	structType, ok := obj.Type().Underlying().(*types.Struct)
	if !ok {
		failErr(fmt.Errorf("type %v is not a struct", obj))
	}

	err := generateBuilder(sourceTypePackage, sourceTypeName, structType)
	if err != nil {
		failErr(err)
	}
}

func loadPackage(path string) *packages.Package {
	cfg := &packages.Config{Mode: packages.NeedTypes | packages.NeedImports}
	pkgs, err := packages.Load(cfg, path)
	if err != nil {
		failErr(fmt.Errorf("loading packages for inspection: %v", err))
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}

	return pkgs[0]
}

func splitSourceType(sourceType string) (string, string) {
	idx := strings.LastIndexByte(sourceType, '.')
	if idx == -1 {
		failErr(fmt.Errorf(`expected qualified type as "pkg/path.MyType"`))
	}
	sourceTypePackage := sourceType[0:idx]
	sourceTypeName := sourceType[idx+1:]
	return sourceTypePackage, sourceTypeName
}

func failErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

func generateBuilder(sourcePackagePath string, sourceTypeName string, structType *types.Struct) error {
	// 1. Get the package of the file with go:generate comment
	goPackage := os.Getenv("GOPACKAGE")

	// 2. Start a new file in this package
	f := NewFile(goPackage)

	// 3. Add a package comment, so IDEs detect files as generated
	f.PackageComment("Code generated by generator, DO NOT EDIT.")

	// 5. Generate changeset type
	builderName := sourceTypeName + "Builder"
	typeLowerCase := firstToLower(sourceTypeName)
	f.Type().Id(builderName).Struct(
		Id(typeLowerCase).Id("*" + sourceTypeName),
	)

	f.Func().
		Id("New" + builderName).
		Params().
		Id("*" + builderName).
		Block(
			Return(Op("&").Id(builderName).Op("{}")),
		)

	for i := 0; i < structType.NumFields(); i++ {
		field := structType.Field(i)
		fieldNameLowerCase := firstToLower(field.Name())

		code := Id(fieldNameLowerCase)
		switch v := field.Type().(type) {
		case *types.Basic:
			code.Id(v.String())
		case *types.Pointer:
			switch ut := v.Elem().(type) {
			case *types.Basic:
				code.Op("*").Id(ut.String())
			case *types.Named:
				typeName := ut.Obj()
				if typeName.Pkg().Path() == sourcePackagePath {
					code.Op("*").Id(typeName.Name())
				} else {
					code.Op("*").Qual(
						typeName.Pkg().Path(),
						typeName.Name(),
					)
				}
			}
		case *types.Named:
			typeName := v.Obj()
			if typeName.Pkg().Path() == sourcePackagePath {
				code.Id(typeName.Name())
			} else {
				code.Qual(
					typeName.Pkg().Path(),
					typeName.Name(),
				)
			}
		default:
			return fmt.Errorf("struct field type not hanled: %T", v)
		}

		f.Func().
			Params(
				Id("b").Id("*"+builderName),
			).
			Id("With"+field.Name()).
			Params(
				code,
			).Id("*"+builderName).
			Block(
				Id("b").Dot(typeLowerCase).Dot(field.Name()).Op("=").Id(fieldNameLowerCase),
				Return(Id("b")),
			)
	}

	f.Func().
		Params(
			Id("b").Id("*" + builderName),
		).
		Id("Build").
		Params().
		Id("*" + sourceTypeName).
		Block(
			Return(Id("b").Dot(typeLowerCase)),
		)

	targetFilename := strings.ToLower(sourceTypeName) + "_builder_gen.go"

	// 7. Write generated file
	return f.Save(targetFilename)
}

func firstToLower(s string) string {
	r, size := utf8.DecodeRuneInString(s)
	if r == utf8.RuneError && size <= 1 {
		return s
	}
	lc := unicode.ToLower(r)
	if r == lc {
		return s
	}
	return string(lc) + s[size:]
}