import json from django.contrib.messages.storage.base import BaseStorage from django.contrib.messages.storage.cookie import MessageEncoder, MessageDecoder from django.utils import six class SessionStorage(BaseStorage): """ Stores messages in the session (that is, django.contrib.sessions). """ session_key = '_messages' def __init__(self, request, *args, **kwargs): assert hasattr(request, 'session'), "The session-based temporary "\ "message storage requires session middleware to be installed, "\ "and come before the message middleware in the "\ "MIDDLEWARE_CLASSES list." super(SessionStorage, self).__init__(request, *args, **kwargs) def _get(self, *args, **kwargs): """ Retrieves a list of messages from the request's session. This storage always stores everything it is given, so return True for the all_retrieved flag. """ return self.deserialize_messages(self.request.session.get(self.session_key)), True def _store(self, messages, response, *args, **kwargs): """ Stores a list of messages to the request's session. """ if messages: self.request.session[self.session_key] = self.serialize_messages(messages) else: self.request.session.pop(self.session_key, None) return [] def serialize_messages(self, messages): encoder = MessageEncoder(separators=(',', ':')) return encoder.encode(messages) def deserialize_messages(self, data): if data and isinstance(data, six.string_types): return json.loads(data, cls=MessageDecoder) return data