diff --git a/CHANGES.rst b/CHANGES.rst index 4b4a23d..9cf5f05 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -4,6 +4,12 @@ Changelog Here you can see the full list of changes between each SQLAlchemy-Utils release. +0.29.7 (2015-03-01) +^^^^^^^^^^^^^^^^^^^ + +- Added Enum representation support for ChoiceType + + 0.29.6 (2015-02-03) ^^^^^^^^^^^^^^^^^^^ diff --git a/setup.py b/setup.py index fa822fa..bac51a7 100644 --- a/setup.py +++ b/setup.py @@ -45,6 +45,7 @@ extras_require = { 'password': ['passlib >= 1.6, < 2.0'], 'color': ['colour>=0.0.4'], 'ipaddress': ['ipaddr'] if not PY3 else [], + 'enum': ['enum34'] if sys.version_info < (3, 4) else [], 'timezone': ['python-dateutil'], 'url': ['furl >= 0.4.1'], 'encrypted': ['cryptography>=0.6'] diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py index 0d1eb70..8329d2b 100644 --- a/sqlalchemy_utils/__init__.py +++ b/sqlalchemy_utils/__init__.py @@ -87,7 +87,7 @@ from .types import ( from .models import Timestamp -__version__ = '0.29.6' +__version__ = '0.29.7' __all__ = ( diff --git a/sqlalchemy_utils/types/choice.py b/sqlalchemy_utils/types/choice.py index a085250..e93fb45 100644 --- a/sqlalchemy_utils/types/choice.py +++ b/sqlalchemy_utils/types/choice.py @@ -2,6 +2,10 @@ from sqlalchemy import types import six from ..exceptions import ImproperlyConfigured from .scalar_coercible import ScalarCoercible +try: + from enum import Enum +except ImportError: + Enum = None class Choice(object): @@ -29,14 +33,21 @@ class Choice(object): class ChoiceType(types.TypeDecorator, ScalarCoercible): """ - ChoiceType offers way of having fixed set of choices for given column. - Columns with ChoiceTypes are automatically coerced to Choice objects. + 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.4+ (the + enum34_ backported package on PyPI is compatible too for ``< 3.4``). + .. _enum34: https://pypi.python.org/pypi/enum34 + + 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(self.Base): + class User(Base): TYPES = [ (u'admin', u'Admin'), (u'regular-user', u'Regular user') @@ -51,6 +62,25 @@ class ChoiceType(types.TypeDecorator, ScalarCoercible): user = User(type=u'admin') user.type # Choice(type='admin', value=u'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 # ChoiceType is very useful when the rendered values change based on user's @@ -61,7 +91,7 @@ class ChoiceType(types.TypeDecorator, ScalarCoercible): from babel import lazy_gettext as _ - class User(self.Base): + class User(Base): TYPES = [ (u'admin', _(u'Admin')), (u'regular-user', _(u'Regular user')) @@ -77,17 +107,49 @@ class ChoiceType(types.TypeDecorator, ScalarCoercible): user.type # Choice(type='admin', value=u'Admin') print user.type # u'Admin' + + Or:: + + from enum import Enum + from babel import lazy_gettext as _ + + + class UserType(Enum): + admin = 1 + regular = 2 + + + UserType.admin.label = _(u'Admin') + UserType.regular.label = _(u'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(TYPES)) + + + user = User(type=UserType.admin) + user.type # + + print user.type.label # u'Admin' """ impl = types.Unicode(255) def __init__(self, choices, impl=None): - if not choices: - raise ImproperlyConfigured( - 'ChoiceType needs list of choices defined.' - ) self.choices = choices - self.choices_dict = dict(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 @@ -95,6 +157,26 @@ class ChoiceType(types.TypeDecorator, ScalarCoercible): 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(object): + """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 @@ -111,3 +193,29 @@ class ChoiceType(types.TypeDecorator, ScalarCoercible): if value: return Choice(value, self.choices_dict[value]) return value + + +class EnumTypeImpl(object): + """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): + return self.enum_class(value) if value else None + + def process_bind_param(self, value, dialect): + return self.enum_class(value).value if value else None + + def process_result_value(self, value, dialect): + return self.enum_class(value) if value else None diff --git a/tests/types/test_choice.py b/tests/types/test_choice.py index 50ae141..d82177e 100644 --- a/tests/types/test_choice.py +++ b/tests/types/test_choice.py @@ -1,7 +1,8 @@ from flexmock import flexmock -from pytest import raises +from pytest import raises, mark import sqlalchemy as sa from sqlalchemy_utils import ChoiceType, Choice, ImproperlyConfigured +from sqlalchemy_utils.types.choice import Enum from tests import TestCase @@ -76,3 +77,54 @@ class TestChoiceTypeWithCustomUnderlyingType(TestCase): def test_init_type(self): type_ = ChoiceType([(1, u'something')], impl=sa.Integer) assert type_.impl == sa.Integer + + +@mark.skipif('Enum is None') +class TestEnumType(TestCase): + def create_models(self): + class OrderStatus(Enum): + unpaid = 1 + paid = 2 + + class Order(self.Base): + __tablename__ = 'order' + id_ = sa.Column(sa.Integer, primary_key=True) + status = sa.Column( + ChoiceType(OrderStatus, impl=sa.Integer()), + default=OrderStatus.unpaid + ) + + def __repr__(self): + return 'Order(%r, %r)' % (self.id_, self.status) + + def pay(self): + self.status = OrderStatus.paid + + self.OrderStatus = OrderStatus + self.Order = Order + + def test_parameter_processing(self): + order = self.Order() + + self.session.add(order) + self.session.commit() + + order = self.session.query(self.Order).first() + assert order.status is self.OrderStatus.unpaid + assert order.status.value == 1 + + order.pay() + self.session.commit() + + order = self.session.query(self.Order).first() + assert order.status is self.OrderStatus.paid + assert order.status.value == 2 + + def test_parameter_coercing(self): + order = self.Order() + order.status = 2 + + self.session.add(order) + self.session.commit() + + assert order.status is self.OrderStatus.paid