Best way to do enum in Sqlalchemy?
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.开发者_如何学Pythontype, 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?
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()
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.
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, anamedtuple
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.
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())
For mysql i use its dialect
from sqlalchemy.dialects.mysql import ENUM
...
class Videos(Base):
...
video_type = Column(ENUM('youtube', 'vimeo'))
...
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))
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,
),
)
精彩评论