Django model objects and querysets dehydration/hydration

# Django model objects and querysets dehydration/hydration

# Dehydrates objects that can be dictionaries, lists or tuples containing django
# model objects or django querysets. For each of those, it creates a
# smaller/dehydrated version of it for saving in cache or pickling. The reverse
# operation is also provided so dehydrated objects can also be re-hydrated.

# Example:
# >>> import pickle
# >>> users = list(User.objects.all()[:20])
# >>> print users
# [<User: Indiana Jones>, <User: Bilbo Baggins>, <User: Lara Croft>, <User: Angus MacGyver>, <User: Luke Skywalker>, <User: Obi-Wan Kenobi>, ...]
# >>> pickled_users = pickle.dumps(users)
# >>> print len(pickled_users)
# 17546
# >>> dehydrated_users = dehydrate(users)
# >>> pickled_dehydrated_users = pickle.dumps(dehydrated_users)
# >>> rehydrated_users = hydrate(pickle.loads(pickled_dehydrated_users))
# >>> print rehydrated_users
# [<User: Indiana Jones>, <User: Bilbo Baggins>, <User: Lara Croft>, <User: Angus MacGyver>, <User: Luke Skywalker>, <User: Obi-Wan Kenobi>, ...]
# >>> print len(pickled_dehydrated_users)
# 1471

import sys
from django.db import models
from django.db.models.query import QuerySet
from django.contrib.contenttypes.models import ContentType

def _walk(obj, fnct, raise_exc=False, delayed=None, maxlevels=10, level=0, context=None):
    if context is None:
        context = {}
    if not obj:
        return obj
    if maxlevels and level >= maxlevels:
        return obj
    objid = id(obj)
    if objid in context:
        return obj
    typ = type(obj)
    if typ is str:
        return obj
    if issubclass(typ, dict):
        if objid in context:
            return obj
        context[objid] = True
        ret = {}
        for k, v in obj.items():
            ret[k] = _walk(v, fnct, raise_exc, delayed, maxlevels, level + 1, context)
        del context[objid]
        return ret
    if issubclass(typ, list) or issubclass(typ, tuple) or issubclass(typ, set):
        context[objid] = True
        ret = []
        for o in obj:
            ret.append(_walk(o, fnct, raise_exc, delayed, maxlevels, level + 1, context))
        del context[objid]
        if issubclass(typ, tuple):
            return tuple(ret)
        if issubclass(typ, set):
            return set(ret)
        return ret
    return fnct(obj, raise_exc, delayed)

class Dehydrated(object):
    def __init__(self, obj):
        self._obj = obj

    def __repr__(self):
        try:
            u = unicode(self)
        except (UnicodeEncodeError, UnicodeDecodeError):
            u = '[Bad Unicode data]'
        return str(u'<%s: %s>' % (self.__class__.__name__, u))

    def __str__(self):
        if hasattr(self, '__unicode__'):
            return unicode(self).encode('utf-8')
        return '%s object' % (self.__class__.__name__,)

    def __getstate__(self):
        ret = self.__dict__.copy()
        ret.pop('_obj', None)
        return ret

class DehydratedModel(Dehydrated):
    def __init__(self, obj):
        super(DehydratedModel, self).__init__(obj)
        self.app_label, self.object_name = obj._meta.app_label, obj._meta.object_name
        self.pk_attrname = obj._meta.pk.attname
        setattr(self, self.pk_attrname, getattr(obj, '_pk', obj.pk))
        if self.pk:
            fields = [f.name for f in obj._meta.fields]
            self.data = dict((k, v) for k, v in obj.__dict__.items() if not k.startswith('_') and not k.endswith('_id') and k not in fields)
        else:
            self.data = dict((k, v) for k, v in obj.__dict__.items() if not k.startswith('_'))

    def __str__(self):
        if hasattr(self, '__unicode__'):
            return unicode(self).encode('utf-8')
        return '%s.%s.%s' % (self.app_label, self.object_name, self.pk)

    def __eq__(self, other):
        return isinstance(other, self.__class__) and self.pk == other.pk

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self.pk)

    def _get_pk_val(self, meta=None):
        return getattr(self, self.pk_attrname)

    def _set_pk_val(self, value):
        return setattr(self, self.pk_attrname, value)

    pk = property(_get_pk_val, _set_pk_val)

    @staticmethod
    def delayed_hydrate(raise_exc, delayed):
        for Model, objects in delayed.items():
            for _obj in Model._default_manager.filter(pk__in=objects.keys()):
                obj, dobjs = objects[_obj.pk]
                obj.__dict__.update(_obj.__dict__)
                for dobj in dobjs:
                    dobj._obj = obj
            for pk, objs in objects.items():
                if not objs[0].pk:
                    msg = "%s matching query does not exist. Cannot hydrate model instance %s.%s.%s." % (
                        Model._meta.object_name,
                        Model._meta.app_label,
                        Model._meta.object_name,
                        pk,
                    )
                    if raise_exc:
                        raise Model.DoesNotExist(msg)
                    sys.stderr.write("%s\n" % msg)

    def hydrate(self, raise_exc=False, delayed=None):
        if hasattr(self, '_obj'):
            obj = self._obj
        else:
            content_type = ContentType.objects.get_by_natural_key(self.app_label, self.object_name.lower())
            Model = content_type.model_class()
            if self.pk:
                if delayed is not None:
                    delayed.setdefault(Model, {})
                    if self.pk in delayed[Model]:
                        obj, dobjs = delayed[Model][self.pk]
                        dobjs.append(self)
                    else:
                        obj = Model()
                        obj._pk = self.pk
                        delayed[Model][self.pk] = (obj, [self])
                else:
                    obj = Model._default_manager.get(pk=self.pk)
                    self._obj = obj
            else:
                obj = Model()
                self._obj = obj
            if hasattr(self, 'data'):
                obj.__dict__.update(self.data)
        return obj

class DehydratedQuerySet(Dehydrated):
    def __init__(self, qs):
        super(DehydratedQuerySet, self).__init__(qs)
        self.app_label, self.object_name = qs.model._meta.app_label, qs.model._meta.object_name
        self.query = qs.query

    @staticmethod
    def delayed_hydrate(raise_exc, delayed):
        pass

    def hydrate(self, raise_exc=False, delayed=None):
        if not hasattr(self, '_obj'):
            content_type = ContentType.objects.get_by_natural_key(self.app_label, self.object_name.lower())
            Model = content_type.model_class()
            qs = Model._default_manager.all()
            qs.query = self.query
            self._obj = qs
        return self._obj

def _dehydrate(obj, raise_exc, delayed):
    typ = type(obj)
    if issubclass(typ, models.Model):
        return DehydratedModel(obj)
    if issubclass(typ, QuerySet):
        return DehydratedQuerySet(obj)
    return obj

def _hydrate(obj, raise_exc, delayed):
    typ = type(obj)
    if issubclass(typ, DehydratedModel):
        if delayed is not None:
            delayed.setdefault(DehydratedModel, {})
            delayed = delayed[DehydratedModel]
        return obj.hydrate(raise_exc, delayed)
    if issubclass(typ, DehydratedQuerySet):
        if delayed is not None:
            delayed.setdefault(DehydratedQuerySet, {})
            delayed = delayed[DehydratedQuerySet]
        return obj.hydrate(raise_exc, delayed)
    return obj

def dehydrate(obj, raise_exc=False):
    """
    Dehydrates objects containing django model objects and querysets.
    """
    return _walk(obj, _dehydrate, raise_exc)

def hydrate(obj, raise_exc=False):
    """
    Hydrates objects containing dehydrated django model objects and querysets.
    """
    delayed = {}
    ret = _walk(obj, _hydrate, raise_exc, delayed)
    for klass, _delayed in delayed.items():
        klass.delayed_hydrate(raise_exc, _delayed)
    return ret