openstack/deb-python-sqlalchemy-utils
Code Issues Proposed changes

make ChoiceType be compatible with Enum type.

Browse Source
This commit is contained in:
Jiangge Zhang
2015-01-18 18:54:06 +08:00
parent e6d0e680dd
commit d98956f5a7
7 changed files with 165 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
docs/data_types.rst
View File

@@ -45,13 +45,6 @@ EncryptedType
.. autoclass:: EncryptedType
EnumType
^^^^^^^^
.. module:: sqlalchemy_utils.types.enum
.. autoclass:: EnumType
JSONType
^^^^^^^^

2
sqlalchemy_utils/__init__.py
View File

@@ -65,7 +65,6 @@ from .types import (
DateTimeRangeType,
EmailType,
EncryptedType,
EnumType,
instrumented_list,
InstrumentedList,
IntRangeType,
@@ -145,7 +144,6 @@ __all__ = (
DateTimeRangeType,
EmailType,
EncryptedType,
EnumType,
ExpressionParser,
ImproperlyConfigured,
InstrumentedList,

2
sqlalchemy_utils/types/__init__.py
View File

@@ -6,7 +6,6 @@ from .color import ColorType
from .country import CountryType, Country
from .email import EmailType
from .encrypted import EncryptedType
from .enum import EnumType
from .ip_address import IPAddressType
from .json import JSONType
from .locale import LocaleType
@@ -37,7 +36,6 @@ __all__ = (
DateTimeRangeType,
EmailType,
EncryptedType,
EnumType,
IntRangeType,
IPAddressType,
JSONType,

123
sqlalchemy_utils/types/choice.py
View File

@@ -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,46 @@ 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 +154,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 +190,27 @@ 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

70
sqlalchemy_utils/types/enum.py
View File

@@ -1,70 +0,0 @@
from __future__ import absolute_import
try:
from enum import Enum
except ImportError:
Enum = None
from sqlalchemy import types
from sqlalchemy_utils.exceptions import ImproperlyConfigured
from .scalar_coercible import ScalarCoercible
class EnumType(types.TypeDecorator, ScalarCoercible):
"""
EnumType offers way of integrating with :mod:`enum` in the standard
library of Python 3.4+ or the enum34_ backported package on PyPI.
.. _enum34: https://pypi.python.org/pypi/enum34
::
from enum import Enum
from sqlalchemy_utils import EnumType
class OrderStatus(Enum):
unpaid = 1
paid = 2
class Order(Base):
__tablename__ = 'order'
id = sa.Column(sa.Integer, autoincrement=True)
status = sa.Column(EnumType(OrderStatus))
order = Order()
order.status = OrderStatus.unpaid
session.add(order)
session.commit()
assert user.status is OrderStatus.unpaid
assert user.status.value == 1
assert user.status.name == 'paid'
"""
impl = types.Integer
def __init__(self, enum_class, impl=None, *args, **kwargs):
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.")
super(EnumType, self).__init__(*args, **kwargs)
self.enum_class = enum_class
if impl is not None:
self.impl = types.Integer
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
def _coerce(self, value):
return self.enum_class(value) if value else None

53
tests/types/test_choice.py
View File

@@ -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,53 @@ 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

53
tests/types/test_enum.py
View File

@@ -1,53 +0,0 @@
from pytest import mark
import sqlalchemy as sa
from sqlalchemy_utils.types import enum
from tests import TestCase
@mark.skipif('enum.Enum is None')
class TestEnumType(TestCase):
def create_models(self):
class OrderStatus(enum.Enum):
unpaid = 1
paid = 2
class Order(self.Base):
__tablename__ = 'document'
id_ = sa.Column(sa.Integer, primary_key=True)
status = sa.Column(
enum.EnumType(OrderStatus), 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