Code generation in Go - constructors
A code generation is an interesting concept, where a program generates another program. It’s very popular in the Java world, for instance, tools like Dagger rely heavily on it. In the Go ecosystem, the code generation hasn’t gained a lot of attention so far. The list of code generators can be found on the official repository https://github.com/golang/go/wiki/GoGenerateTools
Building a generator is quite a straight forward task, I wrote https://github.com/domsu/goconstruct to prove that. It’s a simple tool, that generates constructors for strut types. IDEs such as Goland can already generate constructors, so it’s mainly for educational purposes.
I started by looking for go files in the specified directory and processing each file in the processFile function.
fileNames := getGoFileNamesInDirectory(flag.Args()[0])
if len(fileNames) == 0 {
log.Println("No files to process")
return
}
for _, fileName := range fileNames {
processFile(fileName, typeFilter)
}
The go file is parsed into an AST by calling parser.ParseFile which returns ast.File node.
func processFile(fileName string, typeFilter []string) {
fSet := token.NewFileSet()
node, err := parser.ParseFile(fSet, fileName, nil, parser.ParseComments)
if err != nil {
log.Fatal(err)
}
(...)
}
The ast.File represents a parsed content of the go source file, giving the access to all declarations, documentation, imports etc.
type File struct {
Doc *CommentGroup // associated documentation; or nil
Package token.Pos // position of "package" keyword
Name *Ident // package name
Decls []Decl // top-level declarations; or nil
Scope *Scope // package scope (this file only)
Imports []*ImportSpec // imports in this file
Unresolved []*Ident // unresolved identifiers in this file
Comments []*CommentGroup // list of all comments in the source file
}
At this point, my goal comes down to analyzing the ast.File node.
Looking at the template of the output file, it’s clear that I’ve three main problems to solve: getting a package name, getting correct imports, generating constructors
const genFileContent = `// Code generated by goconstruct. DO NOT EDIT
package %s
import (
%s
)
%s`
Getting the package name
The generated file is located in the same directory as the source file. I simply use node.Name.Name to get the package name.
Getting imports
The file being processed most likely contains a list of imports used in functions, variables, structs etc. The output file though, must only contain a list of imports used in struct types. There should be no missing or unused imports.
It’s done by utilising ast.Inspect(node Node, f func(Node) bool), it traverses an AST in depth-first order.
As it traverses the AST, it calls func(Node) bool function notifying about the current node.
The nodes are declared in the ast package, for instance there are nodes such as StructType, FuncType, ForStmt.
The following code inspects the nodes, visiting only struct types and their content, returning a map of package names used in struct fields.
I use the inspectingStruct and depth variables to check whether or not the program is currently inspecting struct’s content.
func getPackageNamesUsedInStructFields(node *ast.File, structTypeSpecs []*ast.TypeSpec) map[string]bool {
result := make(map[string]bool)
var inspectingStruct = false
var depth = 0
ast.Inspect(node, func(n ast.Node) bool {
var shouldInspect = false
if n == nil && inspectingStruct {
depth--
if depth == 0 {
inspectingStruct = false
}
}
if inspectingStruct {
shouldInspect = true
} else if genDecl, genDeclOk := n.(*ast.GenDecl); genDeclOk {
if typeSpec, typeSpecOk := genDecl.Specs[0].(*ast.TypeSpec); typeSpecOk {
for _, v := range structTypeSpecs {
if typeSpec == v {
shouldInspect = true
inspectingStruct = true
break
}
}
}
} else if _, fileOk := n.(*ast.File); fileOk {
shouldInspect = true
}
if selExpr, selExprOk := n.(*ast.SelectorExpr); selExprOk {
if ident, identOk := selExpr.X.(*ast.Ident); identOk {
result[ident.Name] = true
}
}
if n != nil && inspectingStruct {
depth++
}
return shouldInspect
})
return result
}
Generating constructors
A similar approach is used for the constructor generation.
For each *ast.TypeSpec in structTypeSpecs there is one constructor to generate.
I choose the right function prefix - New for an exported struct, new otherwise.
func generateConstructors(fSet *token.FileSet, structTypeSpecs []*ast.TypeSpec) []string {
var constructors []string
for _, structTypeSpec := range structTypeSpecs {
structName := structTypeSpec.Name.Name
var exported = structName[0] == strings.ToUpper(structName)[0]
var functionPrefix = "new"
if exported {
functionPrefix = "New"
}
(...)
}
*ast.TypeSpec contains all data I need to generate constructor’s arguments and the body.
func generateConstructors(fSet *token.FileSet, structTypeSpecs []*ast.TypeSpec) []string {
(...)
var args []string
var body []string
body = append(body, fmt.Sprintf("\ts := %s{}", structName))
structType := structTypeSpec.Type.(*ast.StructType)
for _, field := range structType.Fields.List {
var buf bytes.Buffer
if err := printer.Fprint(&buf, fSet, field.Type); err != nil {
log.Fatal(err)
}
for _, name := range field.Names {
fieldArg := fmt.Sprintf("%s %s", name, buf.String())
args = append(args, fieldArg)
body = append(body, fmt.Sprintf("\ts.%s = %s", name, name))
}
}
(...)
}
Having the function name, arguments and the body is all I need to generate the constructor.
func generateConstructors(fSet *token.FileSet, structTypeSpecs []*ast.TypeSpec) []string {
(...)
constructorStructName := strings.ToUpper(structName)[0:1] + structName[1:]
constructor := fmt.Sprintf("func %s%s(%s) *%s {\n%s\n}\n", functionPrefix, constructorStructName, strings.Join(args, ","), structName, strings.Join(body, "\n"))
constructors = append(constructors, constructor)
}
return constructors
}
The output file
Getting back to the processFile function. I simply use the genFileContent template, combining the package name, imports and constructors - creating the final file.
func processFile(fileName string, typeFilter []string) {
(...)
constructors := generateConstructors(fSet, structTypeSpecsToProcess)
imports := generateImports(node, structTypeSpecsToProcess)
ext := path.Ext(fileName)
outFile := fileName[0:len(fileName)-len(ext)] + "_gen.go"
genFile, err := os.OpenFile(outFile, os.O_RDWR|os.O_CREATE, 0666)
if err != nil {
log.Fatal(err)
}
defer genFile.Close()
if _, err = genFile.WriteString(fmt.Sprintf(genFileContent, node.Name.Name, strings.Join(imports, "\n"), strings.Join(constructors, "\n"))); err != nil {
log.Fatal(err)
}
}
Conclusion
This only touches a surface of a code generation with AST traversal. Go has a pretty good support for building code generators. I didn’t have to use any external tools and I was surprised how straightforward it was.
Hopefully my article along with the library is a good blueprint to help others building more advanced code generators. Do you see in your projects blocks of code that could be generated?