Utiliser Enum avec SQLAlchemy

Type Enum pratique Cependant, si vous souhaitez l'utiliser dans une application Web, vous souhaitez l'utiliser tel quel dans O / R Mapper.

SQLAlchemy facilite la définition de vos propres types qui encapsulent les types existants en utilisant TypeDecorator. Essayez d'encapsuler Integer et d'enregistrer 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)

Nous avons maintenant un type de colonne dédié pour stocker l'enum spécifié par l'argument mot-clé ʻenum_class`. Utilisons-le.

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()

Le résultat de l'exécution sera comme ça

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}]

De cette manière, les valeurs extraites de la base de données sont correctement renvoyées à Enum, et les valeurs entières non-Enum sont maintenant expulsées avec Insert.