serialization.py

import datetime
import decimal
import functools
import re

from dateutil import rrule
from django.utils import simplejson

from project.utils.dt import utils as dt_utils
from project.utils.misc import is_iterator

WEEKDAY_RE = re.compile(r'^weekday\:([0123456])$')

class JSONEncoder(simplejson.JSONEncoder):
    """Custom simplejson.JSONEncoder subclass that supports serialisation of
       datetime.datetime objects (converts them to RFC 2822 compliant strings),
       among other things."""
    def default(self, o):
        if isinstance(o, datetime.datetime):
            return dt_utils.rfc2822_from_datetime(o)
        elif isinstance(o, datetime.date):
            return dt_utils.rfc2822_from_date(o)
        elif isinstance(o,  datetime.time):
            return dt_utils.rfc2822_from_time(o)
        elif isinstance(o, rrule.weekday):
            return 'weekday:%d' % o.weekday
        elif isinstance(o, decimal.Decimal):
            return float(o)
        else:
            return super(JSONEncoder, self).default(o)

class JSONDecoder(simplejson.JSONDecoder):
    def decode_object(self, o):
        """Hook for de-serialising the stuff JSONEncoder encodes.
           In the case of a sequence being passed, this function calls itself recursively
           on each of its items."""
        if isinstance(o, basestring) and dt_utils.RFC_2822_DATETIME_RE.match(o):
            # Is a datetime object
            return dt_utils.datetime_from_rfc2822(o)
        elif isinstance(o, basestring) and dt_utils.RFC_2822_DATE_RE.match(o):
            # Is a date object
            return dt_utils.date_from_rfc2822(o)
        elif isinstance(o, basestring) and dt_utils.RFC_2822_TIME_RE.match(o):
            # Is a time object
            return dt_utils.time_from_rfc2822(o)
        elif isinstance(o, basestring) and WEEKDAY_RE.match(o):
            # Is a weekday object
            index = int(WEEKDAY_RE.match(o).group(1))
            return rrule.weekdays[index]
        elif isinstance(o, float):
            # Floats go to decimals
            return decimal.Decimal(str(o))
        elif isinstance(o, dict):
            # Is a dictionary, check each of its values
            for key, value in o.items():
                o[key] = self.decode_object(value)
            return o
        elif is_iterator(o):
            # If is an iterator, check all of its values and then convert
            # back to the start type
            result = []
            for item in o:
                result.append(self.decode_object(item))
            result = type(o)(result)
            return result
        else:
            # Otherwise, do nothing
            return o
    
    def raw_decode(self, *args, **kwargs):
        decoded = super(JSONDecoder, self).raw_decode(*args, **kwargs)
        decoded = (self.decode_object(decoded[0]), decoded[1])
        return decoded

def _wrapper(func, cls):
    """Decorator for various simplejson functions that take a cls argument, allowing
       it to be defaulted to something different."""
    @functools.wraps(func)
    def inner_wrapper(obj, **kwargs):
        kwargs.setdefault('cls', cls)
        return func(obj, **kwargs)
    return inner_wrapper

dump = _wrapper(func=simplejson.dump, cls=JSONEncoder)
load = _wrapper(func=simplejson.load, cls=JSONDecoder)
dumps = _wrapper(func=simplejson.dumps, cls=JSONEncoder)
loads = _wrapper(func=simplejson.loads, cls=JSONDecoder)