codegen: add ability to insert comments and blank lines into the generated code

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

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)
+}

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{
+									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{
+												List: []ast.Stmt{panicStmt},
-												&ast.ExprStmt{
-													X: &ast.CallExpr{
-														Fun:  ast.NewIdent("panic"),
-														Args: []ast.Expr{ast.NewIdent("err")},
-													},
-												},
-											},
-										},
 											},
+										}
+									}(),
 								},
 							)
 						}(),
@@ -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{