generate.go

package go_constructor

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"reflect"
	"strings"
	"text/template"

	"github.com/hori-ryota/go-codegen/util/gocodeutil"
	"github.com/hori-ryota/go-codegen/util/typeutil"
	"github.com/hori-ryota/go-strcase"
	"golang.org/x/tools/go/loader"
)

const (
	commentMarker = "//genconstructor"
	pointerOpts   = "-p"
)

func Generate(pkgInfo *loader.PackageInfo) (string, error) {
	allImported := typeutil.AllImported(pkgInfo)
	pkgPathToName := make(map[string]string, len(allImported))
	for _, imported := range allImported {
		pkgPathToName[imported.Path()] = imported.Name()
	}

	printer := typeutil.NewPrinter(pkgInfo.Pkg)
	printer.SetPkgPathToName(pkgPathToName)
	b := new(bytes.Buffer)
	for _, file := range pkgInfo.Files {
		for _, decl := range file.Decls {
			decl, ok := decl.(*ast.GenDecl)
			if !ok {
				continue
			}
			if decl.Tok != token.TYPE {
				continue
			}
			for _, spec := range decl.Specs {
				spec := spec.(*ast.TypeSpec)

				docs := make([]*ast.Comment, 0, 10)
				if spec.Doc != nil {
					docs = append(docs, spec.Doc.List...)
				}
				if decl.Doc != nil {
					docs = append(docs, decl.Doc.List...)
				}
				if len(docs) == 0 {
					continue
				}
				hasMarker := false
				hasPointerOpts := false
				for _, comment := range docs {
					if strings.HasPrefix(strings.TrimSpace(comment.Text), commentMarker) {
						hasMarker = true
						for _, s := range strings.Fields(comment.Text) {
							if s == pointerOpts {
								hasPointerOpts = true
								break
							}
						}
						break
					}
				}
				if !hasMarker {
					continue
				}

				obj, ok := pkgInfo.Defs[spec.Name]
				if !ok {
					return "", fmt.Errorf("definition not found for '%s'", spec.Name)
				}
				named := obj.Type().(*types.Named)

				if err := template.Must(template.New("constructor").Funcs(map[string]interface{}{
					"ToLowerCamel":      strcase.ToLowerCamel,
					"PrintRelativeType": printer.PrintRelativeType,
				}).Parse(`
					func New{{.Struct.Obj.Name}}(
							{{- range .Fields}}
								{{- if not .ConstValue}}
									{{ToLowerCamel .Name}} {{PrintRelativeType .Type}},
								{{- end}}
							{{- end}}
						) {{if .IsPointer}}*{{end}}{{PrintRelativeType .Struct}} {
							return {{if .IsPointer}}&{{end}}{{PrintRelativeType .Struct}}{
								{{- range .Fields}}
									{{- if .ConstValue}}
										{{.Name}}: {{.ConstValue}},
									{{- else }}
										{{.Name}}: {{ToLowerCamel .Name}},
									{{- end}}
								{{- end}}
							}
						}
					`)).Execute(b, map[string]interface{}{
					"Struct":    named,
					"Fields":    parseFields(named),
					"IsPointer": hasPointerOpts,
					"printer":   printer,
				}); err != nil {
					return "", err
				}
			}
		}
	}
	if b.Len() == 0 {
		return "", nil
	}

	out := new(bytes.Buffer)

	err := template.Must(template.New("out").Parse(`
			// Code generated by go-codegen go_constructor; DO NOT EDIT.

			package {{ .PackageName }}

			{{ .ImportPackages }}

			{{ .Body }}
		`)).Execute(out, map[string]string{
		"PackageName":    pkgInfo.Pkg.Name(),
		"ImportPackages": typeutil.FmtImports(allImported, pkgInfo.Pkg),
		"Body":           b.String(),
	})
	if err != nil {
		return "", err
	}

	dst, err := gocodeutil.FormatGoFileFromString(out.String())
	if err != nil {
		return out.String(), err
	}
	return dst, nil
}

type Field struct {
	Type       types.Type
	Name       string
	ConstValue string
}

func parseFields(s types.Type) []Field {
	fields := typeutil.TypeToFields(s)
	tags := typeutil.TypeToTags(s)
	dst := make([]Field, 0, len(fields))
	for i := range fields {
		tag := reflect.StructTag(strings.Trim(tags[i], "`"))

		constValue, hasTag := tag.Lookup("required")
		if !hasTag {
			continue
		}

		dst = append(dst, Field{
			Type:       fields[i].Type(),
			Name:       fields[i].Name(),
			ConstValue: constValue,
		})
	}
	return dst
}