main.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"strings"
	"text/template"
	"unicode"
)

const (
	appName     = "go-getter"
	tplSource   = `func (r {{ if .PointerReceiver }}*{{ end}}{{ .StructName }}) {{ .Getter }}() {{ .Return }} { return r.{{ .Field }} }`
	tplHelpText = `
NAME:
    {{ .AppName }} - generate getter methods for struct fields.

DESCRIPTION:
    {{ .AppName }} is a tool to generate getter methods for struct fields.

Usage: {{ .AppName }} [OPTIONS]

OPTIONS:
    -h, -help show help.
    -typ string the source struct type name.
    -src string source file.
    -out string output file.
    -ptr bool use pointer receiver, default is true.

EXAMPLE:
    {{ .AppName }} -typ User -src user.go -out user_getter.go
`
)

var (
	getterTpl = template.Must(template.New("getter").Parse(tplSource))
	helperTpl = template.Must(template.New("helper").Parse(tplHelpText))
)

func main() {
	if err := Run(os.Args[1:]); err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "go-getter: failed: %+v", err)
		os.Exit(1)
	}
}

func Run(args []string) error {
	f := flag.NewFlagSet(appName, flag.ExitOnError)
	f.Usage = func() {
		_ = helperTpl.Execute(os.Stderr, struct{ AppName string }{appName})
	}

	var (
		useptr            bool
		typname, out, src string
	)

	f.StringVar(&typname, "typ", "", "the source struct type name")
	f.StringVar(&src, "src", "", "source file")
	f.StringVar(&out, "out", "", "output file")
	f.BoolVar(&useptr, "ptr", true, "use pointer receiver, default is true")
	if err := f.Parse(args); err != nil {
		return fmt.Errorf("parse flags: %v", err)
	}

	outfile, err := os.Create(out)
	if err != nil {
		return fmt.Errorf("open out file %q: %w", out, err)
	}
	defer outfile.Close()

	fileset := token.NewFileSet()
	astroot, err := parser.ParseFile(fileset, src, nil, parser.DeclarationErrors)
	if err != nil {
		return fmt.Errorf("parse source file: %w", err)
	}

	typ, err := findStructByName(astroot, typname)
	if err != nil {
		return err
	}

	var (
		fields          = make([]Field, 0, len(typ.Fields.List))
		includedImports = make(map[string]struct{})
	)

	for _, field := range typ.Fields.List {
		if len(field.Names) >= 1 && !isFieldIgnored(field) {
			name := field.Names[0].Name
			// excludes exported fields.
			if unicode.IsLower(rune(name[0])) {
				typ := expr2str(field.Type)
				fields = append(fields, Field{Name: name, Type: typ})

				// extract package selector and add into import set.
				if idx := strings.IndexByte(typ, '.'); idx > 0 {
					includedImports[typ[:idx]] = struct{}{}
				}
			}
		}
	}

	var gen bytes.Buffer
	if err := errors.Join(
		writeHeader(&gen, astroot.Name.Name, astroot.Imports, includedImports),
		writeGeneratedGetter(&gen, fields, typname, useptr),
	); err != nil {
		return fmt.Errorf("generate code: %w", err)
	}

	formatted, err := format.Source(gen.Bytes())
	if err != nil {
		return fmt.Errorf("format generated: %w", err)
	}

	_, err = outfile.Write(formatted)
	return err
}

func writeHeader(buf *bytes.Buffer, pkgName string, imports []*ast.ImportSpec, includedImports map[string]struct{}) error {
	_, err1 := fmt.Fprintf(buf, "// Code generated by %s. DO NOT EDIT.\n\n", appName)
	_, err2 := fmt.Fprintf(buf, "package %s\n\n", pkgName)
	if err := errors.Join(err1, err2); err != nil {
		return fmt.Errorf("writing header: %w", err)
	}

	var closeImports bool
	for i, imp := range imports {
		if i == 0 {
			buf.WriteString("import (")
			closeImports = true
		}
		buf.WriteRune('\n')
		impPath := imp.Path.Value
		if imp.Name != nil { // using import with alias.
			aliasImport := imp.Name.Name
			if _, ok := includedImports[aliasImport]; ok {
				buf.WriteString(aliasImport)
				buf.WriteString(" ")
				buf.WriteString(impPath)
			}
		} else { // normal import.
			unquotedSelector := filepath.Base(impPath[1 : len(impPath)-1])
			if _, ok := includedImports[unquotedSelector]; ok {
				buf.WriteString(impPath)
			}
		}
	}

	if closeImports {
		buf.WriteString(")\n\n")
	}

	return nil
}

func writeGeneratedGetter(buf *bytes.Buffer, fields []Field, typname string, useptr bool) error {
	for _, f := range fields {
		err := getterTpl.Execute(buf, struct {
			StructName      string
			Getter          string
			Return          string
			Field           string
			PointerReceiver bool
		}{
			StructName:      typname,
			Getter:          strings.ToUpper(f.Name[:1]) + f.Name[1:],
			Return:          f.Type,
			Field:           f.Name,
			PointerReceiver: useptr,
		})
		if err != nil {
			return fmt.Errorf("excute template for field: %+v, err: %w", f, err)
		}
		buf.WriteRune('\n')
	}
	return nil
}

func findStructByName(node ast.Node, name string) (*ast.StructType, error) {
	var spec *ast.TypeSpec
	ast.Inspect(node, func(v ast.Node) bool {
		if t, ok := v.(*ast.TypeSpec); ok {
			spec = t
		}
		return spec == nil
	})

	if spec == nil {
		return nil, fmt.Errorf("findStructByName: spec %q is not found", name)
	}

	t, ok := spec.Type.(*ast.StructType)
	if !ok {
		return nil, fmt.Errorf("findStructByName: %q is not a struct type", name)
	}

	return t, nil
}

type Field struct {
	Name string
	Type string
}

func isFieldIgnored(f *ast.Field) bool {
	if f.Tag != nil {
		tag := reflect.StructTag(f.Tag.Value[1 : len(f.Tag.Value)-1])
		value, ok := tag.Lookup(appName)
		return ok && slices.Contains(strings.Split(value, ","), "ignored")
	}
	return false
}

func expr2str(expr ast.Expr) string {
	switch t := expr.(type) {
	default:
		return ""
	case *ast.BasicLit:
		return t.Value
	case *ast.Ident:
		return t.Name
	case *ast.SelectorExpr:
		return fmt.Sprintf("%s.%s", t.X, t.Sel)
	case *ast.StarExpr:
		return "*" + expr2str(t.X)
	case *ast.ArrayType:
		l, e := expr2str(t.Len), expr2str(t.Elt)
		return fmt.Sprintf("[%s]%s", l, e)
	case *ast.IndexExpr: // generic type.
		x, idx := expr2str(t.X), expr2str(t.Index)
		return fmt.Sprintf("%s[%s]", x, idx)
	case *ast.FuncType:
		return funcExpr2str(expr)
	case *ast.MapType:
		key, val := expr2str(t.Key), expr2str(t.Value)
		return fmt.Sprintf("map[%s]%s", key, val)
	case *ast.ChanType:
		val := expr2str(t.Value)
		if t.Dir == ast.SEND {
			return fmt.Sprintf("chan<- %s", val)
		}
		if t.Dir == ast.RECV {
			return fmt.Sprintf("<-chan %s", val)
		}
		return fmt.Sprintf("chan %s", val)
	}
}

func funcExpr2str(expr ast.Expr) string {
	t, ok := expr.(*ast.FuncType)
	if !ok {
		return "<unexpected func expr>"
	}

	params := make([]string, 0, len(t.Params.List))
	for _, p := range t.Params.List {
		params = append(params, expr2str(p.Type))
	}

	results := make([]string, 0, len(t.Results.List))
	for _, r := range t.Results.List {
		results = append(results, expr2str(r.Type))
	}

	ps := strings.Join(params, ", ")
	if len(results) == 0 {
		return fmt.Sprintf("func(%s)", ps)
	}

	rs := strings.Join(results, ", ")
	if len(results) == 1 {
		return fmt.Sprintf("func(%s) %s", ps, rs)
	}

	return fmt.Sprintf("func(%s) (%s)", ps, rs)
}