SQLAlchemy Hybrid Property wrapper to preserve timezones

timestamp.py

#!/usr/bin/env python

from sqlalchemy.ext.hybrid import hybrid_property
from aniso8601.timezone import UTCOffset


def timestamp_with_timezone(timestamp_attr, timezone_attr):
    """Wraps a ``DateTime`` and ``Interval`` column pair to create datetime
    storage that preserves the original timezone information.
    """

    if not timestamp_attr.key:
        raise AttributeError('timestamp_attr must have Column.key set!')
    if not timezone_attr.key:
        raise AttributeError('timezone_attr must have Column.key set!')

    def fget(obj):
        # Fix timezone
        timestamp = getattr(obj, timestamp_attr.key)
        timezone = getattr(obj, timezone_attr.key)
        if timezone is not None:
            tzinfo = UTCOffset()
            tzinfo.setutcdelta(timezone)
            return timestamp.astimezone(tzinfo)
        return timestamp

    def fset(obj, new_timestamp):
        timestamp = getattr(obj, timestamp_attr.key)
        timezone = getattr(obj, timezone_attr.key)
        if new_timestamp.tzinfo is not None:
            timezone = new_timestamp.tzinfo.utcoffset(None)
        else:
            timezone = None
        setattr(obj, timestamp_attr.key, new_timestamp)
        setattr(obj, timezone_attr.key, timezone)

    def expr(cls):
        return timestamp_attr

    return hybrid_property(fget=fget, fset=fset, expr=expr)


class Event(db.Model):
    event_date_timestamp = db.Column(db.DateTime(timezone=True),
                                     key='event_date_timestamp')
    event_date_utcoffset = db.Column(db.Interval,
                                     key='event_date_timezone')

    event_date = timestamp_with_timezone(event_date_timestamp,
                                         event_date_utcoffset)

# Note that the timezone is preserved:
#
#     >>> tzinfo = UTCOffset()
#     >>> tzinfo.setutcdelta(timedelta(hours=3))
#     >>> m = Event()
#     >>> m.event_date = datetime.now.replace(tzinfo=tzinfo)
#     >>> m
#     datetime.datetime(2015, 5, 12, 14, 14, 21, 513186, tzinfo=+9:00:00 UTC)
#
# Querying/sorting the date property works as expected:
#
#     >>> q = Event.query.order_by(Event.event_date.desc())
#     >>> print(q)
#     SELECT id, event_date_timestamp, event_date_utcoffset
#     from event
#     ORDER BY event_date_timestamp DESC;