Reconciling String and Enums in Python

I’ve noticed an unfortunate behaviour when working with Flask APIs and SQLAlchemy models which contain Enum fields. I’ll explain it with an example, and then we’ll write some code to work around it.

Say we have a model Event which represents a given life event like a concert, meeting, or whatever, containing a day Enum field.

class Event(db.Model):
    __tablename__ = 'events'

    class Day(Enum):
        MONDAY = 1
        TUESDAY = 2
        WEDNESDAY = 3
        THURSDAY = 4
        FRIDAY = 5
        SATURDAY = 6
        SUNDAY = 7

    day = db.Column(db.Enum(Day, name='event_days'))

And an API endpoint in our Flask app to update an Event:

@api.route('/events/<int:eid>', methods=['PUT'])
def update_event(eid):
    '''
    Update an Event
    '''
    payload = request.get_json()
    event = Event.query.get(eid)

    # ...

    return jsonify(data=event)

Say we’re sent a PUT /events/1 request with the following payload:

{
    "day": "TUESDAY"
}

In this case, the payload variable in the update_event() function will be the dict representation of the JSON. This means our day will be of type str, not Enum.

This causes some complications if we want to compare the value received before actually updating and committing the change to the Event. This is even worse if we haven’t overridden the __str__ Enum function which means str(Event.Day.MONDAY) returns "Day.MONDAY" and not "MONDAY".

In other words, we have to always do:

if payload['day'] == Event.Day.MONDAY.name:

or:

if payload['day'] in str(Event.Day.MONDAY):

And what if we’re reusing a function which sometimes receives the committed, Enum version, and sometimes receives the str version?

With more complicated control statements, this issue will bloat our code with unnecessary characters. Here’s a real-life example of this that I experienced while overriding the __eq__ function for an Address model:

if self.address_type and other.address_type:
    if not (str(self.address_type) in str(other.address_type)
            or str(other.address_type) in str(self.address_type)):
        return False

Not a pretty sight. So what can we do?

The solution

When something annoys me more than a few times while coding, it’s worth investing some time to come up with a clean workaround which will reduce cognitive load and lines of code.

This is how I solved the issue, and what you could do (I’m sure there are other ways around this, and would love to hear from you if you have a better way!).

First, create a CommonEnum class in your models/base.py (it’s a good idea to inherit from a common class anyways as you might need to extend the Enum default behaviour in other cases too).

class CommonEnum(Enum):
    '''An extended Enum class to act as a base for all models'''

Then, just override the equality function with this:

class CommonEnum(Enum):
    '''An extended Enum class to act as a base for all models'''

    __hash__ = Enum.__hash__

    def __eq__(self, other):
        if isinstance(other, str):
            return self.name == other
        return super().__eq__(other)

This means that if we’re comparing an Enum to a String, it will compare the .name of the Enum, which is a String. Otherwise, it will run the default Enum equality function.

Simple, short, and now we can remove .name or str() everywhere, and stop worrying about casting depending on the context.

I know this helped me a lot. I hope it helps you too!