Best way to do enum in Sqlalchemy?

Question:

I’m reading about sqlalchemy and I saw following code:

employees_table = Table('employees', metadata,
    Column('employee_id', Integer, primary_key=True),
    Column('name', String(50)),
    Column('manager_data', String(50)),
    Column('engineer_info', String(50)),
    Column('type', String(20), nullable=False)
)

employee_mapper = mapper(Employee, employees_table, 
    polymorphic_on=employees_table.c.type, polymorphic_identity='employee')
manager_mapper = mapper(Manager, inherits=employee_mapper, polymorphic_identity='manager')
engineer_mapper = mapper(Engineer, inherits=employee_mapper, polymorphic_identity='engineer')

Should I make ‘type’ an int, with constants in a library? Or should I make just make type an enum?

Asked By: Timmy

||

Answers:

SQLAlchemy has an Enum type since 0.6:
http://docs.sqlalchemy.org/en/latest/core/type_basics.html?highlight=enum#sqlalchemy.types.Enum

Although I would only recommend its usage if your database has a native enum type. Otherwise I would personally just use an int.

Answered By: Wolph

Python’s enumerated types are directly acceptable by the SQLAlchemy Enum type as of SQLAlchemy 1.1:

import enum
from sqlalchemy import Integer, Enum

class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3

class MyClass(Base):
    __tablename__ = 'some_table'
    id = Column(Integer, primary_key=True)
    value = Column(Enum(MyEnum))

Note that above, the string values “one”, “two”, “three” are persisted, not the integer values.

For older versions of SQLAlchemy, I wrote a post which creates its own Enumerated type (http://techspot.zzzeek.org/2011/01/14/the-enum-recipe/)

from sqlalchemy.types import SchemaType, TypeDecorator, Enum
from sqlalchemy import __version__
import re

if __version__ < '0.6.5':
    raise NotImplementedError("Version 0.6.5 or higher of SQLAlchemy is required.")

class EnumSymbol(object):
    """Define a fixed symbol tied to a parent class."""

    def __init__(self, cls_, name, value, description):
        self.cls_ = cls_
        self.name = name
        self.value = value
        self.description = description

    def __reduce__(self):
        """Allow unpickling to return the symbol 
        linked to the DeclEnum class."""
        return getattr, (self.cls_, self.name)

    def __iter__(self):
        return iter([self.value, self.description])

    def __repr__(self):
        return "<%s>" % self.name

class EnumMeta(type):
    """Generate new DeclEnum classes."""

    def __init__(cls, classname, bases, dict_):
        cls._reg = reg = cls._reg.copy()
        for k, v in dict_.items():
            if isinstance(v, tuple):
                sym = reg[v[0]] = EnumSymbol(cls, k, *v)
                setattr(cls, k, sym)
        return type.__init__(cls, classname, bases, dict_)

    def __iter__(cls):
        return iter(cls._reg.values())

class DeclEnum(object):
    """Declarative enumeration."""

    __metaclass__ = EnumMeta
    _reg = {}

    @classmethod
    def from_string(cls, value):
        try:
            return cls._reg[value]
        except KeyError:
            raise ValueError(
                    "Invalid value for %r: %r" % 
                    (cls.__name__, value)
                )

    @classmethod
    def values(cls):
        return cls._reg.keys()

    @classmethod
    def db_type(cls):
        return DeclEnumType(cls)

class DeclEnumType(SchemaType, TypeDecorator):
    def __init__(self, enum):
        self.enum = enum
        self.impl = Enum(
                        *enum.values(), 
                        name="ck%s" % re.sub(
                                    '([A-Z])', 
                                    lambda m:"_" + m.group(1).lower(), 
                                    enum.__name__)
                    )

    def _set_table(self, table, column):
        self.impl._set_table(table, column)

    def copy(self):
        return DeclEnumType(self.enum)

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

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        return self.enum.from_string(value.strip())

if __name__ == '__main__':
    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy import Column, Integer, String, create_engine
    from sqlalchemy.orm import Session

    Base = declarative_base()

    class EmployeeType(DeclEnum):
        part_time = "P", "Part Time"
        full_time = "F", "Full Time"
        contractor = "C", "Contractor"

    class Employee(Base):
        __tablename__ = 'employee'

        id = Column(Integer, primary_key=True)
        name = Column(String(60), nullable=False)
        type = Column(EmployeeType.db_type())

        def __repr__(self):
             return "Employee(%r, %r)" % (self.name, self.type)

    e = create_engine('sqlite://', echo=True)
    Base.metadata.create_all(e)

    sess = Session(e)

    sess.add_all([
        Employee(name='e1', type=EmployeeType.full_time),
        Employee(name='e2', type=EmployeeType.full_time),
        Employee(name='e3', type=EmployeeType.part_time),
        Employee(name='e4', type=EmployeeType.contractor),
        Employee(name='e5', type=EmployeeType.contractor),
    ])
    sess.commit()

    print sess.query(Employee).filter_by(type=EmployeeType.contractor).all()
Answered By: zzzeek

Note: the following is outdated. You should use sqlalchemy.types.Enum now, as recommended by Wolph. It’s particularly nice as it complies with PEP-435 since SQLAlchemy 1.1.


I like zzzeek’s recipe at http://techspot.zzzeek.org/2011/01/14/the-enum-recipe/, but I changed two things:

  • I’m using the Python name of the EnumSymbol also as the name in the database, instead of using its value. I think that’s less confusing. Having a separate value is still useful, e.g. for creating popup menus in the UI. The description can be considered a longer version of the value that can be used e.g. for tooltips.
  • In the original recipe, the order of the EnumSymbols is arbitrary, both when you iterate over them in Python and also when you do an “order by” on the database. But often I want to have a determinate order. So I changed the order to be alphabetic if you set the attributes as strings or tuples, or the order in which the values are declared if you explicitly set the attributes as EnumSymbols – this is using the same trick as SQLAlchemy does when it orders the Columns in DeclarativeBase classes.

Examples:

class EmployeeType(DeclEnum):
    # order will be alphabetic: contractor, part_time, full_time
    full_time = "Full Time"
    part_time = "Part Time"
    contractor = "Contractor"

class EmployeeType(DeclEnum):
    # order will be as stated: full_time, part_time, contractor
    full_time = EnumSymbol("Full Time")
    part_time = EnumSymbol("Part Time")
    contractor = EnumSymbol("Contractor")

Here is the modified recipe; it uses the OrderedDict class available in Python 2.7:

import re

from sqlalchemy.types import SchemaType, TypeDecorator, Enum
from sqlalchemy.util import set_creation_order, OrderedDict


class EnumSymbol(object):
    """Define a fixed symbol tied to a parent class."""

    def __init__(self, value, description=None):
        self.value = value
        self.description = description
        set_creation_order(self)

    def bind(self, cls, name):
        """Bind symbol to a parent class."""
        self.cls = cls
        self.name = name
        setattr(cls, name, self)

    def __reduce__(self):
        """Allow unpickling to return the symbol linked to the DeclEnum class."""
        return getattr, (self.cls, self.name)

    def __iter__(self):
        return iter([self.value, self.description])

    def __repr__(self):
        return "<%s>" % self.name


class DeclEnumMeta(type):
    """Generate new DeclEnum classes."""

    def __init__(cls, classname, bases, dict_):
        reg = cls._reg = cls._reg.copy()
        for k in sorted(dict_):
            if k.startswith('__'):
                continue
            v = dict_[k]
            if isinstance(v, basestring):
                v = EnumSymbol(v)
            elif isinstance(v, tuple) and len(v) == 2:
                v = EnumSymbol(*v)
            if isinstance(v, EnumSymbol):
                v.bind(cls, k)
                reg[k] = v
        reg.sort(key=lambda k: reg[k]._creation_order)
        return type.__init__(cls, classname, bases, dict_)

    def __iter__(cls):
        return iter(cls._reg.values())


class DeclEnum(object):
    """Declarative enumeration.

    Attributes can be strings (used as values),
    or tuples (used as value, description) or EnumSymbols.
    If strings or tuples are used, order will be alphabetic,
    otherwise order will be as in the declaration.

    """

    __metaclass__ = DeclEnumMeta
    _reg = OrderedDict()

    @classmethod
    def names(cls):
        return cls._reg.keys()

    @classmethod
    def db_type(cls):
        return DeclEnumType(cls)


class DeclEnumType(SchemaType, TypeDecorator):
    """DeclEnum augmented so that it can persist to the database."""

    def __init__(self, enum):
        self.enum = enum
        self.impl = Enum(*enum.names(), name="ck%s" % re.sub(
            '([A-Z])', lambda m: '_' + m.group(1).lower(), enum.__name__))

    def _set_table(self, table, column):
        self.impl._set_table(table, column)

    def copy(self):
        return DeclEnumType(self.enum)

    def process_bind_param(self, value, dialect):
        if isinstance(value, EnumSymbol):
            value = value.name
        return value

    def process_result_value(self, value, dialect):
        if value is not None:
            return getattr(self.enum, value.strip())
Answered By: Cito

I’m not really knowledgeable in SQLAlchemy but this approach by Paulo seemed much simpler to me.
I didn’t need user-friendly descriptions, so I went with it.

Quoting Paulo (I hope he doesn’t mind my reposting it here):

Python’s namedtuple collection to the rescue. As the name implies, a namedtuple is a tuple with each item having a name. Like an ordinary tuple, the items are immutable. Unlike an ordinary tuple, an item’s value can be accessed through its name using the dot notation.

Here is a utility function for creating a namedtuple:

from collections import namedtuple

def create_named_tuple(*values):
     return namedtuple('NamedTuple', values)(*values)

The * before the values variable is for “unpacking” the items of the
list so that each item is passed as an individual argument to the
function.

To create a namedtuple, just invoke the above function with the needed
values:

>>> project_version = create_named_tuple('alpha', 'beta', 'prod')
NamedTuple(alpha='alpha', beta='beta', prod='prod')

We can now use the project_version namedtuple to specify the values of the version field.

class Project(Base):
     ...
     version = Column(Enum(*project_version._asdict().values(), name='projects_version'))
     ...

This works great for me and is so much simpler than the other solutions that I previously found.

Answered By: Dan Abramov

For mysql i use its dialect

from sqlalchemy.dialects.mysql import ENUM

... 

class Videos(Base):
    ...
    video_type  = Column(ENUM('youtube', 'vimeo'))
    ...
Answered By: Ricky Levi

This is a method I use – using IntEnum

from enum import IntEnum
class GenderType(IntEnum):
    FEMALE:         int = 1
    MALE:           int = 2
    TRANSGENDER:    int = 3

class Citizen(Base):
    __tablename__ = 'citizen'

    user_uuid:      int = Column(UUID(as_uuid=True), primary_key=True)
    gender_type:    int = Column(Integer, nullable=False, default=GenderType.MALE)
    full_name:      str = Column(String(64))
    address:        str = Column(String(128))
Answered By: Thinkal VB

This and related StackOverflow thread answers resort to PostgreSQL or other dialect-specific typing. However, generic support may be easily achieved in SQLAlchemy that is also compatible with Alembic migrations.

If the backend doesn’t support Enum, SQLAlchemy and alembic can facilitate enforcing constraints on varchar and similar types to mimic enumerated column types.

First, import the Python enum, the SQLAlchemy Enum, and your SQLAlchemy declarative base wherever you’re going to declare your custom SQLAlchemy Enum column type.

import enum
from sqlalchemy import Enum
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

Let’s take OP’s original Python enumerated class:

class PostStatus(enum.Enum):
    DRAFT='draft'
    APPROVE='approve'
    PUBLISHED='published'

Now we create a SQLAlchemy Enum instantiation:

PostStatusType: Enum = Enum(
    PostStatus,
    name="post_status_type",
    create_constraint=True,
    metadata=Base.metadata,
    validate_strings=True,
)

When you run your Alembic alembic revision --autogenerate -m "Revision Notes" and try to apply the revision with alembic upgrade head, you’ll likely get an error about the type not existing. For example:

...
sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedObject) type "post_status_type" does not exist
LINE 10:  post_status post_status_type NOT NULL,
...

To fix this, import your SQLAlchemy Enum class and add the following to your upgrade() and downgrade() functions in the Alembic autogenerated revision script.

from myproject.database import PostStatusType
...
def upgrade() -> None:
    PostStatusType.create(op.get_bind(), checkfirst=True)
    ... the remainder of the autogen code...
def downgrade() -> None:
    ...the autogen code...
    PostStatusType.drop(op.get_bind(), checkfirst=True)

Finally, be sure to update the auto-generated sa.Column() declaration in the table(s) using the enumerated type to simply reference the SQLAlchemy Enum type instead of using Alembic’s attempt to re-declare it. For example in def upgrade() -> None:

op.create_table(
    "my_table",
    sa.Column(
        "post_status",
        PostStatusType,
        nullable=False,
    ),
)
Answered By: Brent
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.