Refactoring Go switch statements

When writing Go code, I often end up with lots of enums that I use to fork my logic in a switch statement. Take this enum:

type MyEnum int

const (
	One   MyEnum = iota
	Two   MyEnum = iota
	Three MyEnum = iota
	Four  MyEnum = iota
	Five  MyEnum = iota
)
}

I’ll then end up with a switch in a part of my code, like so

switch myEnum {
	case One:
		err := DoSomeOtherStuff()
		if err != nil {
			return err
		}
	case Two:
		err := DoSomeMagicalStuff()
		if err != nil {
			return err
		}
	case Three:
		err := DoSomeExoticStuff()
		if err != nil {
			return err
		}
	case Four:
		err := DoSomeOtherStuff()
		if err != nil {
			return err
		}
	case Five:
		err := DoSomeStuff()
		if err != nil {
			return err
		}
}
}

When the enums are small, with only a few entries in them, this is all rather nice and readable. But take an enum that’s 10, 20, 100 entries long and a switch becomes way too long. The approach I prefer to take in these cases is construct a map containing all the required functions associated with the enum value as the key for the map.

var myMap = map[MyEnum]func() error{
	One:   DoSomeOtherStuff,
	Two:   DoSomeMagicalStuff,
	Three: DoSomeExoticStuff,
	Four:  DoSomeOtherStuff,
	Five:  DoSomeStuff,
}
}

The switch then can be gotten rid of. Instead, it becomes a map lookup:

if myFunc, ok := myMap[myEnum]; ok {
	err := myFunc()
	if err != nil {
		return err
	}
} else {
    // the default case would go here
}
}

This allows for more concise code and easier extension in the future. However, I do not use this everywhere either. If the enum is small and not going to change often(famous last words) I’ll leave the switch in its place. It does not work in places where you have functions with very different signatures for each of the switch cases either.

EDIT: Handling all the errors

As Jonathan Gold states in the comments below, you can also handle all the errors at once, so the initial switch becomes something along the lines of:

var err error
switch myEnum {
	case One:
		err = DoSomeOtherStuff()
	case Two:
		err = DoSomeMagicalStuff()
	case Three:
		err = DoSomeExoticStuff()
	case Four:
		err = DoSomeOtherStuff()
	case Five:
		err = DoSomeStuff()
}
if err != nil {
	return err
}
}

Thanks, Jonathan.

How do you approach your switches? What other tricks do you use? Let me know in the comments below.

Comments