Quick: Strong Enums in Go

In modern programming, enums represent a closed number of values within an existing type. Go does not make this easy to represent. I'd argue that a good enum representation really only has 2 properties:

1. Given a value of the enum type, you know it's one of the supported values.
2. It's easy to work with.

Go does have a standard pattern for enum support, but it violates the first property. Borrowed with modification from gobyexample.com/enums:

package server

type State int

const (
	StateIdle State = iota
	StateConnected
	StateError
	StateRetrying
)

var stateName = map[State]string{
	StateIdle:      "idle",
	StateConnected: "connected",
	StateError:     "error",
	StateRetrying:  "retrying",
}

func (s State) String() string {
	return stateName[s]
}

This representation is pretty easy to work with:

• Creating a new value is as simple as naming it: server.StateIdle.
• fmt.Stringer is implemented correctly.
• fmt.GoStringer should be implemented for good test assertions.
• Marshallable to JSON, but without any validation.

The most important problem is that any int can be cast into a server.State:

server.State(42) // This compiles

We can demonstrate some of these problems with tests:

package server_test

import (
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"

	"example.com/server"
)

func TestCannotMarshalInvalidValue(t *testing.T) {
	_, err := json.Marshal(struct {
		ServerState server.State `json:"state"`
	}{ServerState: server.State(42)})
	assert.Error(t, err)
}

func TestCannotUnmarshalInvalidValue(t *testing.T) {
	var example struct {
		ServerState server.State `json:"state"`
	}
	assert.Error(t, json.Unmarshal([]byte(`{"state": 42}`), &example))
}

Both tests fail with:

	Error:      	An error is expected but got nil.

In TestCannotMarshalInvalidValue, we can cast an arbitrary unchecked value into server.State and then marshal it. TestCannotUnmarshalInvalidValue shows us we can unmarshal an arbitrary unchecked value into server.State. Both are bad, since they violate our first principle: "Given a value of the enum type, you know it's one of the supported values."

We can do better. Here is my preferred pattern for a bullet-proof enum in Go:

package server

type State struct{ value int }

var (
	StateIdle      = State{0} // Default value
	StateConnected = State{1}
	StateError     = State{2}
	StateRetrying  = State{3}
)

func (s State) String() string {
	switch s {
	case StateIdle:
		return "idle"
	case StateConnected:
		return "connected"
	case StateError:
		return "error"
	case StateRetrying:
		return "retrying"
	default:
		panic("impossible")
	}
}

This already prevents external packages from constructing invalid enum values, but it's not super ergonomic to use. To get good error messages in tests, I always implement fmt.GoStringer:

func (s State) GoString() string {
	switch s {
	case StateIdle:
		return "server.StateIdle"
	case StateConnected:
		return "server.StateConnected"
	case StateError:
		return "server.StateError"
	case StateRetrying:
		return "server.StateRetrying"
	default:
		panic("impossible")
	}
}

This affects the error messages generated by %#v (and well behaved tests):

 	Error:      	Not equal: 
-	            	expected: server.State{value:1}
+	            	expected: server.StateConnected
-	            	actual  : server.State{value:0}
+	            	actual  : server.StateIdle

Without implementing any (un)marshalling code, this type isn't really possible to marshal or unmarshal, which is correct but restrictive. If we want to support encoding to JSON, YAML and other well-behaved encoding libraries, we need to implement encoding.TextMarshaler and encoding.TextUnmarshaler:

// Implements encoding.TextMarshaler
func (s State) MarshalText() ([]byte, error) { return []byte(s.String()), nil }

// Implements encoding.TextUnmarshaler
func (s *State) UnmarshalText(text []byte) error {
	switch string(text) {
	case "idle":
		*s = StateIdle
	case "connected":
		*s = StateConnected
	case "error":
		*s = StateError
	case "retrying":
		*s = StateRetrying
	default:
		return fmt.Errorf("%q is an invalid server state", string(text))
	}
	return nil
}

All together, the struct based enum is just as easy to work with as the type wrapper enum, but it's fully type-safe outside its package.