Skip to content

Commit 1f2aefa

Browse files
claudeclaude
committed
feat(poet): add expression types and additional statement types
Add new expression types that implement the Expr interface: - CallExpr: function/method calls - StructLit: struct literals with Multiline option - SliceLit: slice literals - TypeCast: type conversions - FuncLit: anonymous function literals - Selector: field/method selection (a.b.c) - Index: array/slice indexing Add new statement types: - GoStmt: goroutine launch (go f()) - Continue: continue statement with optional label - Break: break statement with optional label Update generator.go to use StructLit for: - New function (compact single-line format) - WithTx method (multi-line format) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
1 parent 42f48db commit 1f2aefa

File tree

3 files changed

+256
-26
lines changed

3 files changed

+256
-26
lines changed

‎internal/codegen/golang/generator.go‎

Lines changed: 43 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -127,14 +127,20 @@ func (g *CodeGenerator) addDBCodeStd(f *poet.File) {
127127
f.Decls = append(f.Decls, poet.Func{
128128
Name: "New",
129129
Results: []poet.Param{{Type: "Queries", Pointer: true}},
130-
Stmts: []poet.Stmt{poet.Return{Values: []string{"&Queries{}"}}},
130+
Stmts: []poet.Stmt{poet.Return{Values: []string{
131+
poet.StructLit{Type: "Queries", Pointer: true}.Render(),
132+
}}},
131133
})
132134
} else {
133135
f.Decls = append(f.Decls, poet.Func{
134136
Name: "New",
135137
Params: []poet.Param{{Name: "db", Type: "DBTX"}},
136138
Results: []poet.Param{{Type: "Queries", Pointer: true}},
137-
Stmts: []poet.Stmt{poet.Return{Values: []string{"&Queries{db: db}"}}},
139+
Stmts: []poet.Stmt{poet.Return{Values: []string{
140+
poet.StructLit{Type: "Queries", Pointer: true, Fields: [][2]string{
141+
{"db", "db"},
142+
}}.Render(),
143+
}}},
138144
})
139145
}
140146

@@ -302,23 +308,21 @@ func (g *CodeGenerator) addDBCodeStd(f *poet.File) {
302308

303309
// WithTx method
304310
if !g.tctx.EmitMethodsWithDBArgument {
305-
var withTxBody strings.Builder
306-
withTxBody.WriteString("\treturn &Queries{\n")
307-
withTxBody.WriteString("\t\tdb: tx,\n")
311+
withTxFields := [][2]string{{"db", "tx"}}
308312
if g.tctx.EmitPreparedQueries {
309-
withTxBody.WriteString("\t\ttx: tx,\n")
313+
withTxFields = append(withTxFields, [2]string{"tx", "tx"})
310314
for _, query := range g.tctx.GoQueries {
311-
fmt.Fprintf(&withTxBody, "\t\t%s: q.%s,\n", query.FieldName, query.FieldName)
315+
withTxFields = append(withTxFields, [2]string{query.FieldName, "q." + query.FieldName})
312316
}
313317
}
314-
withTxBody.WriteString("\t}\n")
315-
316318
f.Decls = append(f.Decls, poet.Func{
317319
Recv: &poet.Param{Name: "q", Type: "*Queries"},
318320
Name: "WithTx",
319321
Params: []poet.Param{{Name: "tx", Type: "*sql.Tx"}},
320322
Results: []poet.Param{{Type: "*Queries"}},
321-
Stmts: []poet.Stmt{poet.RawStmt{Code: withTxBody.String()}},
323+
Stmts: []poet.Stmt{poet.Return{Values: []string{
324+
poet.StructLit{Type: "Queries", Pointer: true, Multiline: true, Fields: withTxFields}.Render(),
325+
}}},
322326
})
323327
}
324328
}
@@ -354,14 +358,20 @@ func (g *CodeGenerator) addDBCodePGX(f *poet.File) {
354358
f.Decls = append(f.Decls, poet.Func{
355359
Name: "New",
356360
Results: []poet.Param{{Type: "Queries", Pointer: true}},
357-
Stmts: []poet.Stmt{poet.Return{Values: []string{"&Queries{}"}}},
361+
Stmts: []poet.Stmt{poet.Return{Values: []string{
362+
poet.StructLit{Type: "Queries", Pointer: true}.Render(),
363+
}}},
358364
})
359365
} else {
360366
f.Decls = append(f.Decls, poet.Func{
361367
Name: "New",
362368
Params: []poet.Param{{Name: "db", Type: "DBTX"}},
363369
Results: []poet.Param{{Type: "Queries", Pointer: true}},
364-
Stmts: []poet.Stmt{poet.Return{Values: []string{"&Queries{db: db}"}}},
370+
Stmts: []poet.Stmt{poet.Return{Values: []string{
371+
poet.StructLit{Type: "Queries", Pointer: true, Fields: [][2]string{
372+
{"db", "db"},
373+
}}.Render(),
374+
}}},
365375
})
366376
}
367377

@@ -382,7 +392,11 @@ func (g *CodeGenerator) addDBCodePGX(f *poet.File) {
382392
Name: "WithTx",
383393
Params: []poet.Param{{Name: "tx", Type: "pgx.Tx"}},
384394
Results: []poet.Param{{Type: "Queries", Pointer: true}},
385-
Stmts: []poet.Stmt{poet.Return{Values: []string{"&Queries{\n\t\tdb: tx,\n\t}"}}},
395+
Stmts: []poet.Stmt{poet.Return{Values: []string{
396+
poet.StructLit{Type: "Queries", Pointer: true, Multiline: true, Fields: [][2]string{
397+
{"db", "tx"},
398+
}}.Render(),
399+
}}},
386400
})
387401
}
388402
}
@@ -409,6 +423,7 @@ func (g *CodeGenerator) addModelsCode(f *poet.File) {
409423
f.Decls = append(f.Decls, poet.ConstBlock{Consts: consts})
410424

411425
// Scan method
426+
typeCast := poet.TypeCast{Type: enum.Name, Value: "s"}.Render()
412427
f.Decls = append(f.Decls, poet.Func{
413428
Recv: &poet.Param{Name: "e", Type: enum.Name, Pointer: true},
414429
Name: "Scan",
@@ -421,24 +436,21 @@ func (g *CodeGenerator) addModelsCode(f *poet.File) {
421436
{
422437
Values: []string{"[]byte"},
423438
Body: []poet.Stmt{
424-
poet.Assign{
425-
Left: []string{"*e"}, Op: "=",
426-
Right: []string{fmt.Sprintf("%s(s)", enum.Name)},
427-
},
439+
poet.Assign{Left: []string{"*e"}, Op: "=", Right: []string{typeCast}},
428440
},
429441
},
430442
{
431443
Values: []string{"string"},
432444
Body: []poet.Stmt{
433-
poet.Assign{
434-
Left: []string{"*e"}, Op: "=",
435-
Right: []string{fmt.Sprintf("%s(s)", enum.Name)},
436-
},
445+
poet.Assign{Left: []string{"*e"}, Op: "=", Right: []string{typeCast}},
437446
},
438447
},
439448
{
440449
Body: []poet.Stmt{poet.Return{Values: []string{
441-
fmt.Sprintf(`fmt.Errorf("unsupported scan type for %s: %%T", src)`, enum.Name),
450+
poet.CallExpr{
451+
Func: "fmt.Errorf",
452+
Args: []string{fmt.Sprintf(`"unsupported scan type for %s: %%T"`, enum.Name), "src"},
453+
}.Render(),
442454
}}},
443455
},
444456
},
@@ -509,7 +521,10 @@ func (g *CodeGenerator) addModelsCode(f *poet.File) {
509521
Cond: "!ns.Valid",
510522
Body: []poet.Stmt{poet.Return{Values: []string{"nil", "nil"}}},
511523
},
512-
poet.Return{Values: []string{fmt.Sprintf("string(ns.%s)", enum.Name), "nil"}},
524+
poet.Return{Values: []string{
525+
poet.TypeCast{Type: "string", Value: "ns." + enum.Name}.Render(),
526+
"nil",
527+
}},
513528
},
514529
})
515530

@@ -537,14 +552,16 @@ func (g *CodeGenerator) addModelsCode(f *poet.File) {
537552

538553
// AllValues method
539554
if g.tctx.EmitAllEnumValues {
540-
var valuesList strings.Builder
555+
var enumValues []string
541556
for _, c := range enum.Constants {
542-
fmt.Fprintf(&valuesList, "\t\t%s,\n", c.Name)
557+
enumValues = append(enumValues, c.Name)
543558
}
544559
f.Decls = append(f.Decls, poet.Func{
545560
Name: fmt.Sprintf("All%sValues", enum.Name),
546561
Results: []poet.Param{{Type: "[]" + enum.Name}},
547-
Stmts: []poet.Stmt{poet.Return{Values: []string{fmt.Sprintf("[]%s{\n%s\t}", enum.Name, valuesList.String())}}},
562+
Stmts: []poet.Stmt{poet.Return{Values: []string{
563+
poet.SliceLit{Type: enum.Name, Values: enumValues}.Render(),
564+
}}},
548565
})
549566
}
550567
}

‎internal/poet/ast.go‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// that properly support comment placement.
33
package poet
44

5+
import "strings"
6+
57
// File represents a Go source file.
68
type File struct {
79
BuildTags string
@@ -220,3 +222,150 @@ type VarDecl struct {
220222
}
221223

222224
func (VarDecl) isStmt() {}
225+
226+
// GoStmt represents a go statement (goroutine).
227+
type GoStmt struct {
228+
Call string // The function call to run as a goroutine
229+
}
230+
231+
func (GoStmt) isStmt() {}
232+
233+
// Continue represents a continue statement.
234+
type Continue struct {
235+
Label string // Optional label
236+
}
237+
238+
func (Continue) isStmt() {}
239+
240+
// Break represents a break statement.
241+
type Break struct {
242+
Label string // Optional label
243+
}
244+
245+
func (Break) isStmt() {}
246+
247+
// Expr is an interface for expression types that can be rendered to strings.
248+
// These can be used in Return.Values, Assign.Right, etc.
249+
type Expr interface {
250+
Render() string
251+
}
252+
253+
// CallExpr represents a function or method call expression.
254+
type CallExpr struct {
255+
Func string // Function name or receiver.method
256+
Args []string // Arguments
257+
}
258+
259+
func (c CallExpr) Render() string {
260+
var b strings.Builder
261+
b.WriteString(c.Func)
262+
b.WriteString("(")
263+
for i, arg := range c.Args {
264+
if i > 0 {
265+
b.WriteString(", ")
266+
}
267+
b.WriteString(arg)
268+
}
269+
b.WriteString(")")
270+
return b.String()
271+
}
272+
273+
// StructLit represents a struct literal expression.
274+
type StructLit struct {
275+
Type string // Type name (e.g., "Queries")
276+
Pointer bool // If true, prefix with &
277+
Multiline bool // If true, always use multi-line format
278+
Fields [][2]string // Field name-value pairs (use slice to preserve order)
279+
}
280+
281+
func (s StructLit) Render() string {
282+
var b strings.Builder
283+
if s.Pointer {
284+
b.WriteString("&")
285+
}
286+
b.WriteString(s.Type)
287+
b.WriteString("{")
288+
if len(s.Fields) <= 2 && !s.Multiline {
289+
// Compact format for small struct literals
290+
for i, f := range s.Fields {
291+
if i > 0 {
292+
b.WriteString(", ")
293+
}
294+
b.WriteString(f[0])
295+
b.WriteString(": ")
296+
b.WriteString(f[1])
297+
}
298+
} else if len(s.Fields) > 0 {
299+
// Multi-line format for larger struct literals or when explicitly requested
300+
b.WriteString("\n")
301+
for _, f := range s.Fields {
302+
b.WriteString("\t\t")
303+
b.WriteString(f[0])
304+
b.WriteString(": ")
305+
b.WriteString(f[1])
306+
b.WriteString(",\n")
307+
}
308+
b.WriteString("\t")
309+
}
310+
b.WriteString("}")
311+
return b.String()
312+
}
313+
314+
// SliceLit represents a slice literal expression.
315+
type SliceLit struct {
316+
Type string // Element type (e.g., "interface{}")
317+
Values []string // Elements
318+
}
319+
320+
func (s SliceLit) Render() string {
321+
var b strings.Builder
322+
b.WriteString("[]")
323+
b.WriteString(s.Type)
324+
b.WriteString("{")
325+
for i, v := range s.Values {
326+
if i > 0 {
327+
b.WriteString(", ")
328+
}
329+
b.WriteString(v)
330+
}
331+
b.WriteString("}")
332+
return b.String()
333+
}
334+
335+
// TypeCast represents a type conversion expression.
336+
type TypeCast struct {
337+
Type string // Target type
338+
Value string // Value to convert
339+
}
340+
341+
func (t TypeCast) Render() string {
342+
return t.Type + "(" + t.Value + ")"
343+
}
344+
345+
// FuncLit represents an anonymous function literal.
346+
type FuncLit struct {
347+
Params []Param
348+
Results []Param
349+
Body []Stmt
350+
}
351+
352+
// Note: FuncLit.Render() is implemented in render.go since it needs renderStmts
353+
354+
// Selector represents a field or method selector expression (a.b.c).
355+
type Selector struct {
356+
Parts []string // e.g., ["r", "rows", "0", "Field"] for r.rows[0].Field
357+
}
358+
359+
func (s Selector) Render() string {
360+
return strings.Join(s.Parts, ".")
361+
}
362+
363+
// Index represents an index or slice expression.
364+
type Index struct {
365+
Expr string // Base expression
366+
Index string // Index value (or "start:end" for slice)
367+
}
368+
369+
func (i Index) Render() string {
370+
return i.Expr + "[" + i.Index + "]"
371+
}

‎internal/poet/render.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,12 @@ func renderStmt(b *strings.Builder, s Stmt, indent string) {
315315
renderCallStmt(b, s, indent)
316316
case VarDecl:
317317
renderVarDecl(b, s, indent)
318+
case GoStmt:
319+
renderGoStmt(b, s, indent)
320+
case Continue:
321+
renderContinue(b, s, indent)
322+
case Break:
323+
renderBreak(b, s, indent)
318324
}
319325
}
320326

@@ -459,3 +465,61 @@ func renderVarDecl(b *strings.Builder, v VarDecl, indent string) {
459465
}
460466
b.WriteString("\n")
461467
}
468+
469+
func renderGoStmt(b *strings.Builder, g GoStmt, indent string) {
470+
b.WriteString(indent)
471+
b.WriteString("go ")
472+
b.WriteString(g.Call)
473+
b.WriteString("\n")
474+
}
475+
476+
func renderContinue(b *strings.Builder, c Continue, indent string) {
477+
b.WriteString(indent)
478+
b.WriteString("continue")
479+
if c.Label != "" {
480+
b.WriteString(" ")
481+
b.WriteString(c.Label)
482+
}
483+
b.WriteString("\n")
484+
}
485+
486+
func renderBreak(b *strings.Builder, br Break, indent string) {
487+
b.WriteString(indent)
488+
b.WriteString("break")
489+
if br.Label != "" {
490+
b.WriteString(" ")
491+
b.WriteString(br.Label)
492+
}
493+
b.WriteString("\n")
494+
}
495+
496+
// RenderFuncLit renders a function literal to a string.
497+
// This is used by FuncLit.Render() and can also be called directly.
498+
func RenderFuncLit(f FuncLit) string {
499+
var b strings.Builder
500+
b.WriteString("func(")
501+
renderParams(&b, f.Params)
502+
b.WriteString(")")
503+
if len(f.Results) > 0 {
504+
b.WriteString(" ")
505+
if len(f.Results) == 1 && f.Results[0].Name == "" {
506+
if f.Results[0].Pointer {
507+
b.WriteString("*")
508+
}
509+
b.WriteString(f.Results[0].Type)
510+
} else {
511+
b.WriteString("(")
512+
renderParams(&b, f.Results)
513+
b.WriteString(")")
514+
}
515+
}
516+
b.WriteString(" {\n")
517+
renderStmts(&b, f.Body, "\t")
518+
b.WriteString("}")
519+
return b.String()
520+
}
521+
522+
// Render implements the Expr interface for FuncLit.
523+
func (f FuncLit) Render() string {
524+
return RenderFuncLit(f)
525+
}

0 commit comments

Comments
 (0)