from enum import Enum
from sqlalchemy import types
from ..exceptions import ImproperlyConfigured
from .scalar_coercible import ScalarCoercible
class Choice:
def __init__(self, code, value):
self.code = code
self.value = value
def __eq__(self, other):
if isinstance(other, Choice):
return self.code == other.code
return other == self.code
def __hash__(self):
return hash(self.code)
def __ne__(self, other):
return not (self == other)
def __str__(self):
return str(self.value)
def __repr__(self):
return 'Choice(code={code}, value={value})'.format(
code=self.code,
value=self.value
)
[docs]class ChoiceType(ScalarCoercible, types.TypeDecorator):
"""
ChoiceType offers way of having fixed set of choices for given column. It
could work with a list of tuple (a collection of key-value pairs), or
integrate with :mod:`enum` in the standard library of Python 3.
Columns with ChoiceTypes are automatically coerced to Choice objects while
a list of tuple been passed to the constructor. If a subclass of
:class:`enum.Enum` is passed, columns will be coerced to :class:`enum.Enum`
objects instead.
::
class User(Base):
TYPES = [
('admin', 'Admin'),
('regular-user', 'Regular user')
]
__tablename__ = 'user'
id = sa.Column(sa.Integer, primary_key=True)
name = sa.Column(sa.Unicode(255))
type = sa.Column(ChoiceType(TYPES))
user = User(type='admin')
user.type # Choice(code='admin', value='Admin')
Or::
import enum
class UserType(enum.Enum):
admin = 1
regular = 2
class User(Base):
__tablename__ = 'user'
id = sa.Column(sa.Integer, primary_key=True)
name = sa.Column(sa.Unicode(255))
type = sa.Column(ChoiceType(UserType, impl=sa.Integer()))
user = User(type=1)
user.type # <UserType.admin: 1>
ChoiceType is very useful when the rendered values change based on user's
locale:
::
from babel import lazy_gettext as _
class User(Base):
TYPES = [
('admin', _('Admin')),
('regular-user', _('Regular user'))
]
__tablename__ = 'user'
id = sa.Column(sa.Integer, primary_key=True)
name = sa.Column(sa.Unicode(255))
type = sa.Column(ChoiceType(TYPES))
user = User(type='admin')
user.type # Choice(code='admin', value='Admin')
print user.type # 'Admin'
Or::
from enum import Enum
from babel import lazy_gettext as _
class UserType(Enum):
admin = 1
regular = 2
UserType.admin.label = _('Admin')
UserType.regular.label = _('Regular user')
class User(Base):
__tablename__ = 'user'
id = sa.Column(sa.Integer, primary_key=True)
name = sa.Column(sa.Unicode(255))
type = sa.Column(ChoiceType(UserType, impl=sa.Integer()))
user = User(type=UserType.admin)
user.type # <UserType.admin: 1>
print user.type.label # 'Admin'
"""
impl = types.Unicode(255)
cache_ok = True
def __init__(self, choices, impl=None):
self.choices = tuple(choices) if isinstance(choices, list) else choices
if (
Enum is not None and
isinstance(choices, type) and
issubclass(choices, Enum)
):
self.type_impl = EnumTypeImpl(enum_class=choices)
else:
self.type_impl = ChoiceTypeImpl(choices=choices)
if impl:
self.impl = impl
@property
def python_type(self):
return self.impl.python_type
def _coerce(self, value):
return self.type_impl._coerce(value)
def process_bind_param(self, value, dialect):
return self.type_impl.process_bind_param(value, dialect)
def process_result_value(self, value, dialect):
return self.type_impl.process_result_value(value, dialect)
class ChoiceTypeImpl:
"""The implementation for the ``Choice`` usage."""
def __init__(self, choices):
if not choices:
raise ImproperlyConfigured(
'ChoiceType needs list of choices defined.'
)
self.choices_dict = dict(choices)
def _coerce(self, value):
if value is None:
return value
if isinstance(value, Choice):
return value
return Choice(value, self.choices_dict[value])
def process_bind_param(self, value, dialect):
if value and isinstance(value, Choice):
return value.code
return value
def process_result_value(self, value, dialect):
if value:
return Choice(value, self.choices_dict[value])
return value
class EnumTypeImpl:
"""The implementation for the ``Enum`` usage."""
def __init__(self, enum_class):
if Enum is None:
raise ImproperlyConfigured(
"'enum34' package is required to use 'EnumType' in Python "
"< 3.4"
)
if not issubclass(enum_class, Enum):
raise ImproperlyConfigured(
"EnumType needs a class of enum defined."
)
self.enum_class = enum_class
def _coerce(self, value):
if value is None:
return None
return self.enum_class(value)
def process_bind_param(self, value, dialect):
if value is None:
return None
return self.enum_class(value).value
def process_result_value(self, value, dialect):
return self._coerce(value)