DEV Community

Viktoras
Viktoras

Posted on • Originally published at dizzy.zone on

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
)
Enter fullscreen mode Exit fullscreen mode

}

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
        }
}
Enter fullscreen mode Exit fullscreen mode

}

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,
}
Enter fullscreen mode Exit fullscreen mode

}

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
}
Enter fullscreen mode Exit fullscreen mode

}

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
}
Enter fullscreen mode Exit fullscreen mode

}

Thanks, Jonathan.

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

Top comments (0)