Skip to content

Commit 2eb3639

Browse files
committed
refactor(poet): add Indent field to FuncLit for nested function literals
- Add Indent field to FuncLit to support configurable body indentation - Update RenderFuncLit to use the Indent field (defaults to "\t") - Convert batchmany innerFunc from manual string building to poet.FuncLit This eliminates the last remaining manual string building with tabs in the batch functions, using structured poet AST throughout. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent 8c25017 commit 2eb3639

File tree

3 files changed

+42
-19
lines changed

3 files changed

+42
-19
lines changed

internal/codegen/golang/generator.go

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2119,23 +2119,37 @@ func (g *CodeGenerator) addBatchCodePGX(f *poet.File) {
21192119
})
21202120

21212121
case ":batchmany":
2122-
// Build inner function literal as string (complex nested structure)
2123-
var innerFunc strings.Builder
2124-
innerFunc.WriteString("func() error {\n")
2125-
innerFunc.WriteString("\t\t\trows, err := b.br.Query()\n")
2126-
innerFunc.WriteString("\t\t\tif err != nil {\n")
2127-
innerFunc.WriteString("\t\t\t\treturn err\n")
2128-
innerFunc.WriteString("\t\t\t}\n")
2129-
innerFunc.WriteString("\t\t\tdefer rows.Close()\n")
2130-
innerFunc.WriteString("\t\t\tfor rows.Next() {\n")
2131-
fmt.Fprintf(&innerFunc, "\t\t\t\tvar %s %s\n", q.Ret.Name, q.Ret.Type())
2132-
fmt.Fprintf(&innerFunc, "\t\t\t\tif err := rows.Scan(%s); err != nil {\n", q.Ret.Scan())
2133-
innerFunc.WriteString("\t\t\t\t\treturn err\n")
2134-
innerFunc.WriteString("\t\t\t\t}\n")
2135-
fmt.Fprintf(&innerFunc, "\t\t\t\titems = append(items, %s)\n", q.Ret.ReturnName())
2136-
innerFunc.WriteString("\t\t\t}\n")
2137-
innerFunc.WriteString("\t\t\treturn rows.Err()\n")
2138-
innerFunc.WriteString("\t\t}()")
2122+
// Build inner function literal body
2123+
var innerFuncBody []poet.Stmt
2124+
innerFuncBody = append(innerFuncBody, poet.Assign{
2125+
Left: []string{"rows", "err"}, Op: ":=", Right: []string{"b.br.Query()"},
2126+
})
2127+
innerFuncBody = append(innerFuncBody, poet.If{
2128+
Cond: "err != nil",
2129+
Body: []poet.Stmt{poet.Return{Values: []string{"err"}}},
2130+
})
2131+
innerFuncBody = append(innerFuncBody, poet.Defer{Call: "rows.Close()"})
2132+
2133+
// Build rows loop body
2134+
var rowsLoopBody []poet.Stmt
2135+
rowsLoopBody = append(rowsLoopBody, poet.VarDecl{Name: q.Ret.Name, Type: q.Ret.Type()})
2136+
rowsLoopBody = append(rowsLoopBody, poet.If{
2137+
Cond: fmt.Sprintf("err := rows.Scan(%s); err != nil", q.Ret.Scan()),
2138+
Body: []poet.Stmt{poet.Return{Values: []string{"err"}}},
2139+
})
2140+
rowsLoopBody = append(rowsLoopBody, poet.Assign{
2141+
Left: []string{"items"}, Op: "=", Right: []string{fmt.Sprintf("append(items, %s)", q.Ret.ReturnName())},
2142+
})
2143+
2144+
innerFuncBody = append(innerFuncBody, poet.For{Cond: "rows.Next()", Body: rowsLoopBody})
2145+
innerFuncBody = append(innerFuncBody, poet.Return{Values: []string{"rows.Err()"}})
2146+
2147+
// Build function literal with proper indentation (3 tabs for body inside for loop inside func)
2148+
innerFunc := poet.FuncLit{
2149+
Results: []poet.Param{{Type: "error"}},
2150+
Body: innerFuncBody,
2151+
Indent: "\t\t\t",
2152+
}
21392153

21402154
// Build main for loop body
21412155
var manyForBody []poet.Stmt
@@ -2157,7 +2171,7 @@ func (g *CodeGenerator) addBatchCodePGX(f *poet.File) {
21572171
},
21582172
})
21592173
manyForBody = append(manyForBody, poet.Assign{
2160-
Left: []string{"err"}, Op: ":=", Right: []string{innerFunc.String()},
2174+
Left: []string{"err"}, Op: ":=", Right: []string{innerFunc.Render() + "()"},
21612175
})
21622176
manyForBody = append(manyForBody, poet.If{
21632177
Cond: "f != nil",

internal/poet/ast.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -360,6 +360,7 @@ type FuncLit struct {
360360
Params []Param
361361
Results []Param
362362
Body []Stmt
363+
Indent string // Base indentation for body statements (default: "\t")
363364
}
364365

365366
// Note: FuncLit.Render() is implemented in render.go since it needs renderStmts

internal/poet/render.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -514,7 +514,15 @@ func RenderFuncLit(f FuncLit) string {
514514
}
515515
}
516516
b.WriteString(" {\n")
517-
renderStmts(&b, f.Body, "\t")
517+
indent := f.Indent
518+
if indent == "" {
519+
indent = "\t"
520+
}
521+
renderStmts(&b, f.Body, indent)
522+
// Write closing brace with one less tab than body content
523+
if len(indent) > 0 {
524+
b.WriteString(indent[:len(indent)-1])
525+
}
518526
b.WriteString("}")
519527
return b.String()
520528
}

0 commit comments

Comments
 (0)