stringer generator

stringer.go

package main

import (
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strings"
)

// EnumInfo holds the extracted enum type name and its constants.
type EnumInfo struct {
	TypeName    string
	ConstList   []string
	PackageName string
}

// parseFile parses the Go file and looks for a type definition, constants, and the package name.
func parseFile(filename string) (*EnumInfo, error) {
	fs := token.NewFileSet()
	node, err := parser.ParseFile(fs, filename, nil, parser.AllErrors)
	if err != nil {
		return nil, err
	}

	var enumInfo EnumInfo
	enumInfo.PackageName = node.Name.Name // Capture the package name

	// Walk the AST to find the type and const declarations
	ast.Inspect(node, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.GenDecl:
			if x.Tok == token.TYPE {
				for _, spec := range x.Specs {
					if typeSpec, ok := spec.(*ast.TypeSpec); ok {
						if ident, ok := typeSpec.Type.(*ast.Ident); ok && ident.Name == "int" {
							enumInfo.TypeName = typeSpec.Name.Name
						}
					}
				}
			} else if x.Tok == token.CONST {
				for _, spec := range x.Specs {
					if vs, ok := spec.(*ast.ValueSpec); ok {
						for _, name := range vs.Names {
							enumInfo.ConstList = append(enumInfo.ConstList, name.Name)
						}
					}
				}
			}
		}
		return true
	})

	return &enumInfo, nil
}

// generateStringMethodAST generates a String() method using AST nodes.
func generateStringMethodAST(enumInfo *EnumInfo) *ast.FuncDecl {
	// Create switch body with cases for each constant
	var cases []ast.Stmt
	for _, constName := range enumInfo.ConstList {
		cases = append(cases, &ast.CaseClause{
			List: []ast.Expr{
				ast.NewIdent(constName),
			},
			Body: []ast.Stmt{
				&ast.ReturnStmt{
					Results: []ast.Expr{
						&ast.BasicLit{
							Kind:  token.STRING,
							Value: `"` + constName + `"`,
						},
					},
				},
			},
		})
	}

	// Add a default case returning "Unknown"
	cases = append(cases, &ast.CaseClause{
		Body: []ast.Stmt{
			&ast.ReturnStmt{
				Results: []ast.Expr{
					&ast.BasicLit{
						Kind:  token.STRING,
						Value: `"Unknown"`,
					},
				},
			},
		},
	})

	// Construct the switch statement
	switchStmt := &ast.SwitchStmt{
		Tag: ast.NewIdent(strings.ToLower(enumInfo.TypeName)),
		Body: &ast.BlockStmt{
			List: cases,
		},
	}

	// Create the function declaration
	funcDecl := &ast.FuncDecl{
		Name: ast.NewIdent("String"),
		Recv: &ast.FieldList{
			List: []*ast.Field{
				{
					Names: []*ast.Ident{ast.NewIdent(strings.ToLower(enumInfo.TypeName))},
					Type:  ast.NewIdent(enumInfo.TypeName),
				},
			},
		},
		Type: &ast.FuncType{
			Params:  &ast.FieldList{},
			Results: &ast.FieldList{List: []*ast.Field{{Type: ast.NewIdent("string")}}},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{switchStmt},
		},
	}

	return funcDecl
}

// writeFile writes the AST to a new file.
func writeFile(filename string, f *ast.File) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	// Format and write the AST to the file
	return format.Node(outFile, token.NewFileSet(), f)
}

func main() {
	inputFile := os.Args[1]
	outputFile := os.Args[2]

	// Parse the input file to get enum information
	enumInfo, err := parseFile(inputFile)
	if err != nil {
		log.Fatalf("Error parsing file: %v", err)
	}

	if enumInfo.TypeName == "" || len(enumInfo.ConstList) == 0 {
		log.Fatalf("No valid enum type and constants found")
	}

	// Generate the String() method using AST nodes
	stringMethod := generateStringMethodAST(enumInfo)

	// Create a new AST file and add the method
	newFile := &ast.File{
		Name:  ast.NewIdent(enumInfo.PackageName), // Set package to match input file
		Decls: []ast.Decl{stringMethod},
	}

	// Write the AST to the output file
	err = writeFile(outputFile, newFile)
	if err != nil {
		log.Fatalf("Error writing to output file: %v", err)
	}

	log.Printf("Generated String() method for type %s in package %s to file %s", enumInfo.TypeName, enumInfo.PackageName, outputFile)
}