Help us understand the problem. What is going on with this article?

SQLAlchemy で Enum を使う

More than 5 years have passed since last update.

便利な 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 で蹴られるようになりました。

Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away