stein.wtf

Emulating enumerated types in Golang

☕️ 6 min read

In this post we look at generating a powerful enumerated type using go generate and abstract syntax tree traversal.

The output of this post is a CLI to generate enumerated types. The full source code is available on github.

Idiomatic Go

Go does not have first class support for enumerated types. One way to define an enumerated type is to construct a set of related constants aided by a type definition. Iota may be used to predefine successive incrementing integer constants. We can define a Color type like so.

package main

import "fmt"

type Color int

const (
    Red Color = iota // 0
    Blue             // 1
)

func main() {
    var b1 Color = Red
    b1 = Red
    fmt.Println(b1) // prints 0

    var b2 Color = 1
    fmt.Println(b2 == Blue) // prints true

    var b3 Color
    b3 = 42
    fmt.Println(b3)  // prints 42
}

This pattern is very common in Go code. Whilst being idiomatic the approach has some drawbacks. There is no static type checking since any integer can be passed as a Color. There is no serialization support - it’s pretty uncommon that a developer will want to serialize this to an integer over the wire or to a database record. There is no support for a readable display value - we’ll need to convert the const to a display value in code.

It’s important to know the idioms of a language and when to break those idioms. More often than not the argument for idioms is used to shut down arguments. This can sometimes be the death of creativity.

Designing an enumerated type

One of Go’s best features is its simplicity - developers coming from other languages can generally be efficient in Go very quickly. On the flip side this imposes constraints, such as the lack of generics which can lead to boilerplate code. To overcome some of these shortcomings the community has embraced code generation as a mechanism to define more powerful and flexible types.

Let’s use this approach to define an enumerated type. One approach would be to generate an enum as a struct. We can also attach methods to the struct. Structs also support a meta tag which will be useful for defining the display value and description.

type ColorEnum struct {
    Red  string `enum:"RED"`
    Blue string `enum:"BLUE"`
}

All we need to do now is to generate an instance of the struct for each field.

var Red  = Color{name: "RED"}
var Blue = Color{name: "BLUE"}

We can then attach methods to the Color struct to support JSON encoding/decoding. We implement the Marshaler interface to support JSON encoding.

func (c Color) MarshalJSON() ([]byte, error) {
    return json.Marshal(c.name)
}

Go will invoke our custom implementation when serializing this type as JSON. Likewise we can implement the Unmarshaler interface which will enable us to consume enum types - this allows us to define enumerated types directly on data transfer objects in our APIs.

func (c *Color) UnmarshalJSON(b []byte) error {
    return json.Unmarshal(b, c.name)
}

We can also attach some helper methods to generate a slice of the display values.

// ColorNames returns the displays values of all enum instances as a slice
func ColorNames() []string { ... }

We’d also like support to generate an enum instance from a string, so let’s add that.

// NewColor generates a new Color from the given display value (name)
func NewColor(value string) (Color, error) { ... }

This is really extensible and you might want to add other methods to return the name, support errors by implement Error() string and support Stringer by implementing String() string.

Generating the code

Traversing the abstract syntax tree

Before rendering a template to generate code we will need to parse the ColorEnum type in the source code. Two common approaches are to use the reflect and ast packages. We will need to scan structs declared at the package level. The ast package has the capability to generate an abstract syntax tree - a traversable data structure representing Go source code. We can then traverse the syntax tree and match a provided type. This type and defined struct tags can then be parsed and used to build a model to generate a template. We start off by loading a go package

cfg := &packages.Config{
    Mode:  packages.LoadSyntax,
    Tests: false,
}
pkgs, err := packages.Load(cfg, patterns...)

The pkgs variable contains syntax trees for each file in the package. The ast.Inspect method can then be used to walk the AST. It takes a function that will be called for each node encountered. We loop over each file and process the syntax tree for that file.

for _, file := range pkg.files {
...
    ast.Inspect(file.file, func(node ast.Node) bool {
        // ...handle node, check if it's something we are interested in
    })
}

The consumer should define this function and then filter by token types they are interested in. You can filter by structs with this check on the node

node.Tok == token.STRUCT { ... }

In our case we filter by struct types that have an enum: tag defined. We simply process each token in the source and construct a model (custom Go struct) based on the data encountered.

Rendering source code

There are a few approaches to generating code. The Stringer tool writes to standard out using the fmt package. Whilst this is easy to get going it becomes unwieldy and difficult to debug as the generator scales. A more sane approach is to use the text/template package and utilise Go’s powerful templating library. It allows you to decouple logic that generates the model from the templating leading to separated concerns and code that is easier to reason about. The resulting type definitions might look like this.

// {{.NewType}} is the enum that instances should be created from
type {{.NewType}} struct {
    name  string
}

// Enum instances
{{- range $e := .Fields}}
var {{.Value}} = {{$.NewType}}{name: "{{.Key}}"}
{{- end}}

... code to generate methods

We can then render the template from our model

t, err := template.New(tmpl).Parse(tmpl)
if err != nil {
    log.Fatal("instance template parse error: ", err)
}

err = t.Execute(buf, model)

It’s best not to worry about formatting when developing the template. The format package has a method which take source code as a parameter and returns formatted Go code, so let Go handle this for you.

func Source(src []byte) ([]byte, error) { ... }

Conclusion

In this post we looked at a method to generate an enumerated type by parsing Go source code. This approach can be used as a template to build other code generators that need to interpret source code. We used Go’s text/template library to render the source code in a maintainable way.

Read the full source code on github.