From d6426bba14e518e2036cb08661c5f9924e8c516a Mon Sep 17 00:00:00 2001 From: wispem-wantex Date: Sun, 15 Feb 2026 14:54:22 -0800 Subject: [PATCH] codegen: add ability to insert comments and blank lines into the generated code --- cmd/subcmd_generate_models.go | 7 +- pkg/codegen/modelgenerate/ast_helpers.go | 171 +++++++++++++++++++- pkg/codegen/modelgenerate/generate_model.go | 43 ++--- 3 files changed, 197 insertions(+), 24 deletions(-) diff --git a/cmd/subcmd_generate_models.go b/cmd/subcmd_generate_models.go index cfdf0d8..2e634b6 100644 --- a/cmd/subcmd_generate_models.go +++ b/cmd/subcmd_generate_models.go @@ -4,7 +4,6 @@ import ( "errors" "fmt" "go/ast" - "go/printer" "go/token" "os" @@ -37,11 +36,9 @@ var generate_model = &cobra.Command{ return ErrNoSuchTable } - fset := token.NewFileSet() - if Must(cmd.Flags().GetBool("test")) { file2 := modelgenerate.GenerateModelTestAST(table, modname) - PanicIf(printer.Fprint(os.Stdout, fset, file2)) + PanicIf(modelgenerate.FprintWithComments(os.Stdout, file2)) } else { decls := []ast.Decl{ &ast.GenDecl{ @@ -81,7 +78,7 @@ var generate_model = &cobra.Command{ Decls: decls, } - PanicIf(printer.Fprint(os.Stdout, fset, file)) + PanicIf(modelgenerate.FprintWithComments(os.Stdout, file)) } return nil diff --git a/pkg/codegen/modelgenerate/ast_helpers.go b/pkg/codegen/modelgenerate/ast_helpers.go index a7ad809..f07d527 100644 --- a/pkg/codegen/modelgenerate/ast_helpers.go +++ b/pkg/codegen/modelgenerate/ast_helpers.go @@ -1,6 +1,30 @@ package modelgenerate -import "go/ast" +import ( + "bytes" + "go/ast" + "go/parser" + "go/printer" + "go/token" + "io" + "sort" +) + +// These just need to be unique (i.e., something that will never be a real function in actual code). +const ( + commentMarker = "__comment__jafkjewfkajwlekfjawlejf" + blankLineMarker = "__blank_line__awefjakwlefjkwalkefj" +) + +// TrailingComments is a side map for attaching end-of-line comments to AST nodes. +// Generators populate this map; FprintWithComments consumes and clears it. +// For a regular statement, the comment appears at the end of the statement's line. +// For a BlockStmt, the comment appears after the closing brace. +// +// This is a terrible implementation (global variable), but Node is an interface (pointer type), +// so there's no risk of cross-contamination really. Still bad for concurrent use, but hopefully +// we don't have to do that. +var TrailingComments = map[ast.Node]string{} // mustCall wraps a call expression in Must(...), producing AST for Must(inner). func mustCall(inner ast.Expr) *ast.CallExpr { @@ -9,3 +33,148 @@ func mustCall(inner ast.Expr) *ast.CallExpr { Args: []ast.Expr{inner}, } } + +// Comment creates a marker statement that will be converted to a real Go comment +// in the generated output by FprintWithComments. +func Comment(text string) ast.Stmt { + return &ast.ExprStmt{X: &ast.CallExpr{ + Fun: ast.NewIdent(commentMarker), + Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: `"` + text + `"`}}, + }} +} + +// BlankLine creates a marker statement that will be converted to a blank line +// in the generated output. The marker occupies a line in the first-pass output; +// when removed, the position gap causes the printer to insert a blank line. +func BlankLine() *ast.ExprStmt { + return &ast.ExprStmt{X: &ast.CallExpr{ + Fun: ast.NewIdent(blankLineMarker), + }} +} + +// FprintWithComments prints an ast.File, converting Comment/BlankLine markers +// into real Go comments and blank lines, preserving Doc comments, and applying +// trailing comments from the TrailingComments side map. +// +// It does a round-trip (print -> parse -> modify -> print) to obtain real token +// positions, which Go's comment system requires. +func FprintWithComments(w io.Writer, file *ast.File) error { + // First pass: print to buffer + // Doc comments on FuncDecl/GenDecl are NOT added to file.Comments here; + // the printer handles them via setComment(d.Doc) when visiting each decl. + var buf bytes.Buffer + fset := token.NewFileSet() + if err := printer.Fprint(&buf, fset, file); err != nil { + return err + } + + // Re-parse to get real positions (ParseComments preserves doc comments) + fset = token.NewFileSet() + parsed, err := parser.ParseFile(fset, "", buf.Bytes(), parser.ParseComments) + if err != nil { + return err + } + + // Convert the tree-of-nodes into a slice-of-nodes + collectNodes := func(node ast.Node) []ast.Node { + var nodes []ast.Node + ast.Inspect(node, func(n ast.Node) bool { + // Filter out comments, as they only appear in the + switch n.(type) { + case *ast.CommentGroup, *ast.Comment: + return false + } + nodes = append(nodes, n) + return true + }) + return nodes + } + + // Parallel walk: apply TrailingComments from the side map. + // Both trees have identical structure (the reparse is just a positioned copy), + // so ast.Inspect visits nodes in the same order. We skip comment nodes to + // avoid mismatches from Doc fields. + if len(TrailingComments) > 0 { + origNodes := collectNodes(file) + reparsedNodes := collectNodes(parsed) + for i, orig := range origNodes { + text, ok := TrailingComments[orig] + if !ok { + continue + } + reparsed := reparsedNodes[i] + parsed.Comments = append(parsed.Comments, &ast.CommentGroup{ + List: []*ast.Comment{{Slash: reparsed.End(), Text: "// " + text}}, + }) + } + } + + extractCommentMarker := func(stmt ast.Stmt) (string, bool) { + expr, ok := stmt.(*ast.ExprStmt) + if !ok { + return "", false + } + call, ok := expr.X.(*ast.CallExpr) + if !ok { + return "", false + } + ident, ok := call.Fun.(*ast.Ident) + if !ok || ident.Name != commentMarker { + return "", false + } + lit := call.Args[0].(*ast.BasicLit) + return lit.Value[1 : len(lit.Value)-1], true + } + + isBlankLineMarker := func(stmt ast.Stmt) bool { + expr, ok := stmt.(*ast.ExprStmt) + if !ok { + return false + } + call, ok := expr.X.(*ast.CallExpr) + if !ok { + return false + } + ident, ok := call.Fun.(*ast.Ident) + return ok && ident.Name == blankLineMarker + } + + // Convert comment and blank-line markers + ast.Inspect(parsed, func(n ast.Node) bool { + // We only care about Block nodes + block, ok := n.(*ast.BlockStmt) + if !ok { + return true + } + + // Check the statements in this block , replacing our artificial Comment(...) and BlankLine(...) nodes + // in the block statement's body with actual ones + filtered := block.List[:0] + for _, stmt := range block.List { + if text, ok := extractCommentMarker(stmt); ok { + // If it's a comment, add it to the fileset's list of Comments + parsed.Comments = append(parsed.Comments, &ast.CommentGroup{ + List: []*ast.Comment{{Slash: stmt.Pos(), Text: "// " + text}}, + }) + } else if isBlankLineMarker(stmt) { + // If it's a blank line, just remove it; the position gap creates a blank line + } else { + // Otherwise: it's a normal statement, so keep it. + filtered = append(filtered, stmt) + } + } + block.List = filtered + return true + }) + + sort.Slice(parsed.Comments, func(i, j int) bool { + return parsed.Comments[i].Pos() < parsed.Comments[j].Pos() + }) + + // Clear side map for next invocation + for k := range TrailingComments { + delete(TrailingComments, k) + } + + return printer.Fprint(w, fset, parsed) +} diff --git a/pkg/codegen/modelgenerate/generate_model.go b/pkg/codegen/modelgenerate/generate_model.go index de31ad6..bb96ee3 100644 --- a/pkg/codegen/modelgenerate/generate_model.go +++ b/pkg/codegen/modelgenerate/generate_model.go @@ -240,7 +240,7 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { List: func() []ast.Stmt { ret := []ast.Stmt{} if hasFks { - ret = append(ret, checkForeignKeyFailuresAssignment) + ret = append(ret, checkForeignKeyFailuresAssignment, BlankLine()) } if hasUpdatedAt { // Auto-timestamps: updated_at @@ -261,7 +261,7 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { // Do create List: append( func() []ast.Stmt { - ret1 := []ast.Stmt{} + ret1 := []ast.Stmt{Comment("Do create")} if hasCreatedAt { // Auto-timestamps: created_at ret1 = append(ret1, &ast.AssignStmt{ @@ -321,23 +321,25 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { }, }, }, - Else: &ast.IfStmt{ - Cond: &ast.BinaryExpr{ - X: ast.NewIdent("err"), - Op: token.NEQ, - Y: ast.NewIdent("nil"), - }, - Body: &ast.BlockStmt{ - List: []ast.Stmt{ - &ast.ExprStmt{ - X: &ast.CallExpr{ - Fun: ast.NewIdent("panic"), - Args: []ast.Expr{ast.NewIdent("err")}, - }, - }, + Else: func() *ast.IfStmt { + panicStmt := &ast.ExprStmt{ + X: &ast.CallExpr{ + Fun: ast.NewIdent("panic"), + Args: []ast.Expr{ast.NewIdent("err")}, }, - }, - }, + } + TrailingComments[panicStmt] = "not a foreign key error" + return &ast.IfStmt{ + Cond: &ast.BinaryExpr{ + X: ast.NewIdent("err"), + Op: token.NEQ, + Y: ast.NewIdent("nil"), + }, + Body: &ast.BlockStmt{ + List: []ast.Stmt{panicStmt}, + }, + } + }(), }, ) }(), @@ -357,6 +359,7 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { Else: &ast.BlockStmt{ // Do update List: []ast.Stmt{ + Comment("Do update"), &ast.AssignStmt{ Lhs: []ast.Expr{ast.NewIdent("result")}, Tok: token.DEFINE, @@ -389,6 +392,10 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { } funcDecl := &ast.FuncDecl{ + Doc: &ast.CommentGroup{List: []*ast.Comment{ + {Text: fmt.Sprintf("// Save%s creates or updates a %s in the database.", tbl.GoTypeName, tbl.GoTypeName)}, + {Text: "// If the item doesn't exist (has no ID set), it will create it; otherwise it will do an update."}, + }}, Recv: dbRecv, Name: ast.NewIdent("Save" + tbl.GoTypeName), Type: &ast.FuncType{