Enum type for SQLAlchemy

enum_type.py

"""
This work is licensed under the MIT License (https://opensource.org/licenses/MIT)

Example:

class Example(Model):
    class Type(enum.Enum):
        one = 1
        two = 2

    type = sa.Column(EnumType(Type), nullable=False)
"""


import sqlalchemy as sa


class EnumType(sa.TypeDecorator):
    impl = sa.SmallInteger

    def __init__(self, data, **kw):
        self.data = data
        super(EnumType, self).__init__(**kw)

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        return value.value

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