← All Articles

Using Python enums in SQLAlchemy models

I'm a big fan of using enums, as they generally force all values appearing in a database column to be in a consistent format. Of course, it's not suited for columns with high cardinality (ie lots of different values) - but if you are only expecting a small set of values AND you are annoyed at values being stored as 'green' vs 'Green' vs 'Grreen' every time an engineer mistypes something, there's lots of reasons to like enums.

This article will show you how to use enums in Python with SQLAlchemy models. Note that I'm not talking about enums at the database level - these are definitely useful but perhaps not as flexible as just defining enums in your model definition (ie changing the former requires a database migration, but the latter only requires a code / definition change). 

You may also be interested in reading about enums from the author of SQLAlchemy, although this article was written in 2011 and may be dated.

So here's how I do it.

Define your enum values

Python enums are special beasts which are essentially key-value stores, ie each enum has both a name and a value attribute.

Here's a sample enum definition:

import enum

class MessageTypes(enum.IntEnum):
    text = 1
    image = 2
    audio = 3
    reply = 4
    unknown = 5

print MessageTypes.text  # <MessageTypes.text: 1>
print MessageTypes.text.name # 'text'
print MessageTypes.text.value # 1
print MessageTypes.text < 1  # False

As you can see, the comparator for an enum uses its value.

Create a custom field type

I prefer to store the enum's values (eg 1 in the example above) rather than its name (eg "text" in the example above).

In order to do so using SQLAlchemy, we have to define a special column type which will automatically store/return the enum's value.

class IntEnum(db.TypeDecorator):
    Enables passing in a Python enum and storing the enum's *value* in the db.
    The default would have stored the enum's *name* (ie the string).
    impl = db.Integer

    def __init__(self, enumtype, *args, **kwargs):
        super(IntEnum, self).__init__(*args, **kwargs)
        self._enumtype = enumtype

    def process_bind_param(self, value, dialect):
        if isinstance(value, int):
            return value

        return value.value

    def process_result_value(self, value, dialect):
        return self._enumtype(value)

Define your SQLALchemy model

Now that this new column type is available, here's how it looks in the model definition.

class Message(BaseModel):

    __tablename__ = 'messages'

    message_type = db.Column(IntEnum(MessageTypes), default=MessageTypes.text)
    ... some other columns ...

This allows you to use the enums in your code, which is very handy since most code editors will autocomplete the set of possible names.

m = Message()
m.message_type = MessageTypes.audio

Be aware though, that this "enum-ness" is enforced by SQLAlchemy and not at the database level. So you can very easily run a raw SQL statement to set the value to anything you like.

Made with JoyBird