Merge branch 'tonyseek-add-enum-support'

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

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

sqlalchemy_utils/__init__.py

@@ -87,7 +87,7 @@ from .types import (
 from .models import Timestamp
 
 
-__version__ = '0.29.6'
+__version__ = '0.29.7'
 
 
 __all__ = (

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  # <UserType.admin: 1>
 
 
     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  # <UserType.admin: 1>
+
+        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

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