Exercise: Code Generator
Difficulty - Advanced
Learning Objectives
- Parse Go source files with go/parser
- Generate Go code programmatically
- Use go/format for code formatting
- Implement go generate integration
- Work with go/template for code generation
- Handle package imports and dependencies
Problem Statement
Create a code generator that reads struct definitions and generates boilerplate code like JSON serialization, validation, and builders.
Core Components
1package codegen
2
3import (
4 "go/ast"
5 "go/parser"
6 "go/token"
7)
8
9type Generator struct {
10 fset *token.FileSet
11 file *ast.File
12 structs []*StructInfo
13}
14
15type StructInfo struct {
16 Name string
17 Fields []*FieldInfo
18}
19
20type FieldInfo struct {
21 Name string
22 Type string
23 Tags string
24}
25
26func New(filename string)
27func Parse() error
28func GenerateJSON()
29func GenerateBuilder()
30func GenerateValidator()
Solution
Click to see the solution
1package codegen
2
3import (
4 "bytes"
5 "fmt"
6 "go/ast"
7 "go/format"
8 "go/parser"
9 "go/token"
10 "strings"
11 "text/template"
12)
13
14type Generator struct {
15 fset *token.FileSet
16 file *ast.File
17 structs []*StructInfo
18 packageName string
19}
20
21type StructInfo struct {
22 Name string
23 Fields []*FieldInfo
24}
25
26type FieldInfo struct {
27 Name string
28 Type string
29 Tags string
30 JSONName string
31}
32
33func New(filename string) {
34 fset := token.NewFileSet()
35 file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
36 if err != nil {
37 return nil, err
38 }
39
40 return &Generator{
41 fset: fset,
42 file: file,
43 structs: make([]*StructInfo, 0),
44 packageName: file.Name.Name,
45 }, nil
46}
47
48func Parse() error {
49 for _, decl := range g.file.Decls {
50 genDecl, ok := decl.(*ast.GenDecl)
51 if !ok || genDecl.Tok != token.TYPE {
52 continue
53 }
54
55 for _, spec := range genDecl.Specs {
56 typeSpec, ok := spec.(*ast.TypeSpec)
57 if !ok {
58 continue
59 }
60
61 structType, ok := typeSpec.Type.(*ast.StructType)
62 if !ok {
63 continue
64 }
65
66 structInfo := &StructInfo{
67 Name: typeSpec.Name.Name,
68 Fields: make([]*FieldInfo, 0),
69 }
70
71 for _, field := range structType.Fields.List {
72 for _, name := range field.Names {
73 fieldInfo := &FieldInfo{
74 Name: name.Name,
75 Type: g.typeToString(field.Type),
76 }
77
78 if field.Tag != nil {
79 fieldInfo.Tags = field.Tag.Value
80 fieldInfo.JSONName = g.extractJSONTag(field.Tag.Value)
81 }
82
83 structInfo.Fields = append(structInfo.Fields, fieldInfo)
84 }
85 }
86
87 g.structs = append(g.structs, structInfo)
88 }
89 }
90
91 return nil
92}
93
94func typeToString(expr ast.Expr) string {
95 switch t := expr.(type) {
96 case *ast.Ident:
97 return t.Name
98 case *ast.SelectorExpr:
99 return fmt.Sprintf("%s.%s", g.typeToString(t.X), t.Sel.Name)
100 case *ast.StarExpr:
101 return "*" + g.typeToString(t.X)
102 case *ast.ArrayType:
103 return "[]" + g.typeToString(t.Elt)
104 case *ast.MapType:
105 return fmt.Sprintf("map[%s]%s",
106 g.typeToString(t.Key), g.typeToString(t.Value))
107 default:
108 return "interface{}"
109 }
110}
111
112func extractJSONTag(tag string) string {
113 tag = strings.Trim(tag, "`")
114 parts := strings.Split(tag, " ")
115 for _, part := range parts {
116 if strings.HasPrefix(part, "json:") {
117 jsonTag := strings.TrimPrefix(part, "json:")
118 jsonTag = strings.Trim(jsonTag, `"`)
119 if comma := strings.Index(jsonTag, ","); comma != -1 {
120 return jsonTag[:comma]
121 }
122 return jsonTag
123 }
124 }
125 return ""
126}
127
128func GenerateJSON() {
129 tmpl := `// Code generated by codegen. DO NOT EDIT.
130
131package {{.Package}}
132
133import "encoding/json"
134
135{{range .Structs}}
136// MarshalJSON implements json.Marshaler
137func MarshalJSON() {
138 type Alias {{.Name}}
139 return json.Marshal(&struct {
140 *Alias
141 }{
142 Alias:(s),
143 })
144}
145
146// UnmarshalJSON implements json.Unmarshaler
147func UnmarshalJSON(data []byte) error {
148 type Alias {{.Name}}
149 aux := &struct {
150 *Alias
151 }{
152 Alias:(s),
153 }
154 return json.Unmarshal(data, aux)
155}
156{{end}}
157`
158
159 return g.executeTemplate(tmpl)
160}
161
162func GenerateBuilder() {
163 tmpl := `// Code generated by codegen. DO NOT EDIT.
164
165package {{.Package}}
166
167{{range .Structs}}
168// {{.Name}}Builder builds {{.Name}}
169type {{.Name}}Builder struct {
170 {{range .Fields}}{{.Name}} {{.Type}}
171 {{end}}
172}
173
174// New{{.Name}}Builder creates a new builder
175func New{{.Name}}Builder() *{{.Name}}Builder {
176 return &{{.Name}}Builder{}
177}
178
179{{range .Fields}}
180// {{.Name}} sets the {{.Name}} field
181func {{.Name}}(v {{.Type}}) *{{.Name}}Builder {
182 b.{{.Name}} = v
183 return b
184}
185{{end}}
186
187// Build constructs the {{.Name}}
188func Build() *{{.Name}} {
189 return &{{.Name}}{
190 {{range .Fields}}{{.Name}}: b.{{.Name}},
191 {{end}}
192 }
193}
194{{end}}
195`
196
197 return g.executeTemplate(tmpl)
198}
199
200func GenerateValidator() {
201 tmpl := `// Code generated by codegen. DO NOT EDIT.
202
203package {{.Package}}
204
205import "errors"
206
207{{range .Structs}}
208// Validate validates {{.Name}}
209func Validate() error {
210 {{range .Fields}}
211 {{if eq .Type "string"}}
212 if s.{{.Name}} == "" {
213 return errors.New("{{.Name}} is required")
214 }
215 {{end}}
216 {{end}}
217 return nil
218}
219{{end}}
220`
221
222 return g.executeTemplate(tmpl)
223}
224
225func executeTemplate(tmplStr string) {
226 tmpl, err := template.New("gen").Parse(tmplStr)
227 if err != nil {
228 return "", err
229 }
230
231 data := struct {
232 Package string
233 Structs []*StructInfo
234 }{
235 Package: g.packageName,
236 Structs: g.structs,
237 }
238
239 var buf bytes.Buffer
240 if err := tmpl.Execute(&buf, data); err != nil {
241 return "", err
242 }
243
244 // Format the generated code
245 formatted, err := format.Source(buf.Bytes())
246 if err != nil {
247 return buf.String(), err // Return unformatted if formatting fails
248 }
249
250 return string(formatted), nil
251}
252
253func WriteFile(filename, content string) error {
254 // Write to file
255 return nil // Implementation omitted for brevity
256}
Usage Example
Define structs in models.go:
1package models
2
3//go:generate go run generator.go models.go
4
5type User struct {
6 ID int `json:"id"`
7 Name string `json:"name"`
8 Email string `json:"email"`
9 Password string `json:"-"`
10}
11
12type Post struct {
13 ID int `json:"id"`
14 Title string `json:"title"`
15 Content string `json:"content"`
16 UserID int `json:"user_id"`
17}
Create generator tool:
1package main
2
3import (
4 "flag"
5 "fmt"
6 "os"
7)
8
9func main() {
10 flag.Parse()
11 if flag.NArg() < 1 {
12 fmt.Println("Usage: generator <file.go>")
13 os.Exit(1)
14 }
15
16 filename := flag.Arg(0)
17 gen, err := codegen.New(filename)
18 if err != nil {
19 fmt.Fprintf(os.Stderr, "Error: %v\n", err)
20 os.Exit(1)
21 }
22
23 if err := gen.Parse(); err != nil {
24 fmt.Fprintf(os.Stderr, "Parse error: %v\n", err)
25 os.Exit(1)
26 }
27
28 // Generate JSON methods
29 jsonCode, err := gen.GenerateJSON()
30 if err != nil {
31 fmt.Fprintf(os.Stderr, "JSON generation error: %v\n", err)
32 os.Exit(1)
33 }
34 fmt.Println(jsonCode)
35
36 // Generate builder pattern
37 builderCode, err := gen.GenerateBuilder()
38 if err != nil {
39 fmt.Fprintf(os.Stderr, "Builder generation error: %v\n", err)
40 os.Exit(1)
41 }
42 fmt.Println(builderCode)
43
44 // Generate validators
45 validatorCode, err := gen.GenerateValidator()
46 if err != nil {
47 fmt.Fprintf(os.Stderr, "Validator generation error: %v\n", err)
48 os.Exit(1)
49 }
50 fmt.Println(validatorCode)
51}
Run code generation:
1go generate ./...
Key Takeaways
- go/parser enables AST-based code generation
- go/format ensures generated code is properly formatted
- text/template simplifies code generation
- go generate integrates generation into build process
- Type information extraction requires careful AST traversal
- Generated code should include build tags and comments
- Consider using go:generate directives in source files
- Validate generated code with go vet and tests