Some Go enum tricks

Hey. This is the first article of this blog. I go over some simple things to warmup my writing skills.

Go is known to be a simple language. It is part of its philosophy, which focuses on productivity and maintainability. But this can sometimes be a problem. I recently came across one of those situations when I needed to use enums. Go does not support enums, but we can approximate them in different ways. The simplest way is to use constants:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
const (
        Sun = iota
        Mercury
        Venus
        Earth
        Mars
        AsteroidBelt
        Jupiter
        Saturn
        Uranus
        Neptune
        Pluto
)

iota is a special keyword which starts at 0 and automatically increments the next constant who belongs to the same const scope. This implementation lacks several properties commonly found in enums. Let’s look at some techniques we can use to improve this.

Here, the enums type default to int. We can define a type alias to int and use this new type to define methods on our enums. This does not provide us with type safety however, meaning that we can pass any int value to a function accepting a type Planet, even if this value is not mapped to a valid planet. Since the default value of an int is 0, we discard this value so that we can differentiate between specific values provided on purpose from the default value of an int which might indicate an error. But how do we make sure that a variable of this type holds a correct planet? We create a method on the Planet type along with a dummy last value _planetNumber that is not exported to automatically check for the correct size. This is useful when we often add or remove constants.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

type Planet int

const (
        _ Planet = iota
        Sun
        Mercury
        Venus
        ...
        _planetNumber
)

func (p *Planet) InTheSolarSystem() bool {
        return p > 0 && p < _planetNumber
}

func main() {

    planets := []Planet{Mercury, Planet(10)}

	for _, planet := range planets {
		if planet.InTheSolarSystem() {
			fmt.Printf("This planet is in the solar system\n")
			continue
		}
		fmt.Printf("This is planet is not in the solar system\n")
	}
}

Nice. Having the enum to hold an int allows us to easily verify if it is valid. Now we might want to know more about it, like its name. Go has a builtin interface called Stringer. Any type that has a method String() string implements this interface. You could manually implement this interface. But if you have a large amount of enums, or are constantly modifying them, manually maintaining a slice of those string representations can be gruesome.

Automated code generation

Like any respectable computer languages, Go is Turing complete. Therefore, it can generate its own code! Using the //go generate some_tool.bin, we can call any binary at compile time. One of those particularly useful binary is stringer: it makes our type automatically implement the Stringer interface by having the name of the constant being its string representation. First run in your terminal:

$ go get golang.org/x/tools/cmd/stringer

In the code, we add a coomented line which is recognized by the go generate command:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

type Planet int

const (
        _ Planet = iota
        Sun
        Mercury
        Venus
        ...
)

//go:generate stringer --type=Planet

Finally, we call go generate in the directory of the package containing the type. To execute all //go:generate statements in your module, you can use:

$ go generate ./...

This will create a file planet_string.go with the following content:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
// Code generated by "stringer --type=Planet"; DO NOT EDIT.

package main

import "strconv"

func _() {
	// An "invalid array index" compiler error signifies that the constant values have changed.
	// Re-run the stringer command to generate them again.
	var x [1]struct{}
	_ = x[Sun-1]
	_ = x[Mercury-2]
	_ = x[Venus-3]
	_ = x[Earth-4]
	_ = x[Mars-5]
	_ = x[Jupiter-6]
	_ = x[Saturn-7]
	_ = x[Uranus-8]
	_ = x[Neptune-9]
	_ = x[Pluto-10]
}

const _Planet_name = "SunMercuryVenusEarthMarsJupiterSaturnUranusNeptunePluto"

var _Planet_index = [...]uint8{0, 3, 10, 15, 20, 24, 31, 37, 43, 50, 55}

func (i Planet) String() string {
	i -= 1
	if i < 0 || i >= Planet(len(_Planet_index)-1) {
		return "Planet(" + strconv.FormatInt(int64(i+1), 10) + ")"
	}
	return _Planet_name[_Planet_index[i]:_Planet_index[i+1]]
}

This is incredibly useful, and I have been using this throughout the Go packages I have been writing in my work. There are many other tools which you can use with go generate, like [ifacemaker]{https://github.com/vburenin/ifacemaker} to generate interfaces from structs and mockery to generate mocks from interfaces.

Reversing the relationship

Now that we can get a string representation of a constant automatically, what about the reverse?

If you look at the file planet_string.go, you can see the array _Planet_index and the string _Planet_name. Those two variables can be used to reverse the relationship and obtain the corresponding Planet type from its string representation !

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

func PlanetFromString(planetName string) Planet {
	index := strings.Index(_Planet_name, planetName)
	for i, j := range _Planet_index{
		if int(j) == index {
			return Planet(i+1)
		}
	}
	return Planet(0)
}

If your constants contain more than 10 “runs” (consecutive integer values), the content of this file will be very different. This is because stringer will use a map _Planet_map instead of the list _Planet_index to balance between space and complexity. In that case, we can reassign this map to a public variable and use it directly to reverse the relationship!

1
var PlanetMap = _Planet_map

That’s all folks.

Nota bene: the Sun and Pluto are no planets.