便利な Enum 型
ですが、 Web アプリで使う場合は O/R Mapper でそのまま使いたいですよね。
SQLAlchemy は、 TypeDecorator
を使って、既存の型をラップした独自型を簡単に定義できます。 Integer をラップして 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)
キーワード引数 enum_class
で指定した Enum を保存する専用のカラム型ができました。使ってみましょう。
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()
実行結果はこうなります
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}]
このように、DBから取得した値が適切に Enum に戻りますし、 Enum じゃない整数値は Insert で蹴られるようになりました。