[PYTHON] Use Enums with SQLAlchemy

Convenient Enum type However, if you want to use it in a web application, you want to use it as it is in O / R Mapper.

SQLAlchemy makes it easy to define your own types that wrap existing types using TypeDecorator. Try wrapping the Integer and saving the Enum.

from __future__ import print_function, division, absolute_import
from sqlalchemy import Integer
from sqlalchemy.types import TypeDecorator


class EnumType(TypeDecorator):
    """Store IntEnum as Integer"""

    impl = Integer

    def __init__(self, *args, **kwargs):
        self.enum_class = kwargs.pop('enum_class')
        TypeDecorator.__init__(self, *args, **kwargs)

    def process_bind_param(self, value, dialect):
        if value is not None:
            if not isinstance(value, self.enum_class):
                raise TypeError("Value should %s type" % self.enum_class)
            return value.value

    def process_result_value(self, value, dialect):
        if value is not None:
            if not isinstance(value, int):
                raise TypeError("value should have int type")
            return self.enum_class(value)

We now have a dedicated column type to store the Enum specified by the keyword argument ʻenum_class`. Let's use it.

import enum
from sqlalchemy import Integer, Column, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker


Base = declarative_base()

Color = enum.Enum("Color", "red green yellow")

class Light(Base):
    __tablename__ = 'some_table'
    id = Column(Integer, primary_key=True)
    color = Column(EnumType(enum_class=Color))


engine = create_engine('sqlite:///:memory:')
Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)
session = Session()

session.add(Light(color=Color.red))
session.commit()

for L in session.query(Light).all():
    print(L.id, L.color)

session.add(Light(color=42))
session.commit()

The execution result will be like this

1 Color.red
Traceback (most recent call last):
...
  File "enum_column_sample.py", line 19, in process_bind_param
    raise TypeError("Value should %s type" % self.enum_class)
sqlalchemy.exc.StatementError: Value should <enum 'Color'> type (original cause: TypeError: Value should <enum 'Color'> type) u'INSERT INTO some_table (color) VALUES (?)' [{'color': 42}]

In this way, the value obtained from the DB is properly returned to the Enum, and the non-Enum integer value is now kicked by Insert.