main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
)

type Struct struct {
	Name        string
	Fields      map[string]string
	InnerFields map[string]string
}

func toTitle(s string) string {
	return strings.ToUpper(s[:1]) + s[1:]
}

func getFieldTypeName(e ast.Expr) string {
	switch ft := e.(type) {
	case *ast.StarExpr:
		return "*" + getFieldTypeName(ft.X)
	case *ast.ArrayType:
		return "[]" + getFieldTypeName(ft.Elt)
	case *ast.MapType:
		return "map[" + getFieldTypeName(ft.Key) + "]" + getFieldTypeName(ft.Value)
	case *ast.SelectorExpr:
		return getFieldTypeName(ft.X) + "." + getFieldTypeName(ft.Sel)
	case *ast.Ident:
		return ft.Name
	case *ast.StructType:
		return "struct{}"
	default:
		fmt.Fprintf(os.Stderr, "unknown type: %+v\n", ft)
		return "unknown"
	}
}

func parseStruct(f *ast.File) *Struct {
	for _, node := range f.Decls {
		genDecl, ok := node.(*ast.GenDecl)
		if !ok {
			continue
		}
		for _, spec := range genDecl.Specs {
			ts, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			st, ok := ts.Type.(*ast.StructType)
			if !ok || (*structName != "" && ts.Name.Name != *structName) {
				continue
			}
			s := &Struct{
				Name:        ts.Name.Name,
				Fields:      map[string]string{},
				InnerFields: map[string]string{},
			}
			for _, field := range st.Fields.List {
				ft := getFieldTypeName(field.Type)
				for _, name := range field.Names {
					if name.Name[0] == strings.ToLower(name.Name)[0] {
						s.InnerFields[name.Name] = ft
					} else {
						s.Fields[name.Name] = ft
					}
				}
			}
			return s
		}
	}
	return nil
}

var (
	fileName   = flag.String("file", "", "file name (default is all)")
	structName = flag.String("struct", "", "struct name (required)")

	structs = []Struct{}
	fset    = token.NewFileSet()
)

func main() {
	flag.Parse()
	if *structName == "" {
		fmt.Fprintln(os.Stderr, "struct name is required")
		return
	}
	pkgs, err := parser.ParseDir(fset, ".", func(fi fs.FileInfo) bool {
		return !(*fileName != "" && fi.Name() != *fileName)
	}, parser.ParseComments)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}
	// Get struct info
	buf := bytes.Buffer{}
	for _, pkg := range pkgs {
		for fn, pf := range pkg.Files {
			buf.Reset()
			buf.WriteString("// Code generated by github.com/virzz/structx. DO NOT EDIT.\n\n")
			buf.WriteString("package " + pkg.Name + "\n\n")
			s := parseStruct(pf)
			if s == nil || s.Fields == nil {
				continue
			}
			// func (s *StructName) WithFieldName(v FieldType) *StructName{
			// 	s.FieldName = v
			// 	return s
			// }
			for name, typ := range s.InnerFields {
				buf.WriteString("func (s *" + s.Name + ") With" + toTitle(name) + "(v " + typ + ") *" + s.Name + "{\n")
				buf.WriteString("s." + name + " = v\n")
				buf.WriteString("return s\n")
				buf.WriteString("}\n")
				buf.WriteString("func (s *" + s.Name + ") " + toTitle(name) + "() " + typ + "{ return s." + name + " }\n")
			}
			buf.WriteString("\n")
			for name, typ := range s.Fields {
				buf.WriteString("func (s *" + s.Name + ") With" + toTitle(name) + "(v " + typ + ") *" + s.Name + "{\n")
				buf.WriteString("s." + name + " = v\n")
				buf.WriteString("return s\n")
				buf.WriteString("}\n")
			}
			genName := strings.ToLower(fn[:len(fn)-len(filepath.Ext(fn))] + "_" + s.Name + ".structx.go")
			// Generate struct methods
			data, err := format.Source(buf.Bytes())
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				return
			}
			err = os.WriteFile(genName, data, 0o644)
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				return
			}
			fmt.Println("Generated struct methods for", s.Name, "to", genName)
		}
	}
}