diff --git a/docs/index.rst b/docs/index.rst index fd69e30..efe5d16 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -191,6 +191,24 @@ or a 16-byte BINARY column or a 32-character CHAR column if not. id = sa.Column(UUIDType(binary=False), primary_key=True) +TimezoneType +------------ + +TimezoneType provides a way for saving timezones (from either the pytz or the dateutil package) objects into database. +TimezoneType saves timezone objects as strings on the way in and converts them back to objects when querying the database. + + +:: + + from sqlalchemy_utils import UUIDType + + class User(Base): + __tablename__ = 'user' + + # Pass backend='pytz' to change it to use pytz (dateutil by default) + timezone = sa.Column(TimezoneType(backend='pytz')) + + API Documentation ----------------- diff --git a/setup.py b/setup.py index 76905af..a6d0171 100644 --- a/setup.py +++ b/setup.py @@ -42,7 +42,8 @@ extras_require = { ], 'password': ['passlib >= 1.6, < 2.0'], 'color': ['colour>=0.0.4'], - 'ipaddress': ['ipaddr'] if not PY3 else [] + 'ipaddress': ['ipaddr'] if not PY3 else [], + 'timezone': ['python-dateutil'] } diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py index fe034b1..f66409b 100644 --- a/sqlalchemy_utils/__init__.py +++ b/sqlalchemy_utils/__init__.py @@ -22,6 +22,7 @@ from .types import ( NumberRangeType, ScalarListType, ScalarListException, + TimezoneType, TSVectorType, UUIDType, ) @@ -58,6 +59,7 @@ __all__ = ( ProxyDict, ScalarListType, ScalarListException, + TimezoneType, TSVectorType, UUIDType, ) diff --git a/sqlalchemy_utils/types/__init__.py b/sqlalchemy_utils/types/__init__.py index b91867a..ac78219 100644 --- a/sqlalchemy_utils/types/__init__.py +++ b/sqlalchemy_utils/types/__init__.py @@ -15,6 +15,7 @@ from .number_range import ( from .password import Password, PasswordType from .phone_number import PhoneNumber, PhoneNumberType from .scalar_list import ScalarListException, ScalarListType +from .timezone import TimezoneType from .uuid import UUIDType @@ -33,6 +34,7 @@ __all__ = ( PhoneNumberType, ScalarListException, ScalarListType, + TimezoneType, UUIDType, ) diff --git a/sqlalchemy_utils/types/timezone.py b/sqlalchemy_utils/types/timezone.py new file mode 100644 index 0000000..df0f6d4 --- /dev/null +++ b/sqlalchemy_utils/types/timezone.py @@ -0,0 +1,74 @@ +import six +from sqlalchemy import types +from sqlalchemy_utils import ImproperlyConfigured + + +class TimezoneType(types.TypeDecorator): + """ + Changes Timezone objects to a string representation on the way in and + changes them back to Timezone objects on the way out. + """ + + impl = types.CHAR(50) + + python_type = None + + def __init__(self, backend='dateutil'): + """ + :param backend: Whether to use 'dateutil' or 'pytz' for timezones. + """ + + self.backend = backend + if backend == 'dateutil': + try: + from dateutil.tz import tzfile + from dateutil.zoneinfo import gettz + + self.python_type = tzfile + self._to = gettz + self._from = lambda x: x._filename + + except ImportError: + raise ImproperlyConfigured( + "'python-dateutil' is required to use the " + "'dateutil' backend for 'TimezoneType'" + ) + + elif backend == 'pytz': + try: + from pytz import tzfile, timezone + + self.python_type = tzfile.DstTzInfo + self._to = timezone + self._from = six.text_type + + except ImportError: + raise ImproperlyConfigured( + "'pytz' is required to use the 'pytz' backend " + "for 'TimezoneType'" + ) + + else: + raise ImproperlyConfigured( + "'pytz' or 'dateutil' are the backends supported for " + "'TimezoneType'" + ) + + def _coerce(self, value): + if value and not isinstance(value, self.python_type): + obj = self._to(value) + if obj is None: + raise ValueError("unknown time zone '%s'" % value) + + return obj + + return value + + def coercion_listener(self, target, value, oldvalue, initiator): + return self._coerce(value) + + def process_bind_param(self, value, dialect): + return self._from(self._coerce(value)) if value else None + + def process_result_value(self, value, dialect): + return self._to(value) if value else None diff --git a/tests/test_timezone.py b/tests/test_timezone.py new file mode 100644 index 0000000..3f6c67d --- /dev/null +++ b/tests/test_timezone.py @@ -0,0 +1,39 @@ +from pytest import mark +import six +import sqlalchemy as sa +from sqlalchemy_utils.types import timezone +from tests import TestCase + + +try: + import dateutil + +except ImportError: + dateutil = None + + +@mark.skipif('dateutil is None') +class TestTimezoneType(TestCase): + def create_models(self): + class Visitor(self.Base): + __tablename__ = 'document' + id = sa.Column(sa.Integer, primary_key=True) + timezone = sa.Column(timezone.TimezoneType) + + def __repr__(self): + return 'Visitor(%r)' % self.id + + self.Visitor = Visitor + + def test_parameter_processing(self): + visitor = self.Visitor( + timezone=u'America/Los_Angeles' + ) + + self.session.add(visitor) + self.session.commit() + + visitor = self.session.query(self.Visitor).filter_by( + timezone='America/Los_Angeles').first() + + assert visitor is not None