from django.contrib.messages.storage.base import BaseStorage


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.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] = messages
        else:
            self.request.session.pop(self.session_key, None)
        return []