session.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import json
  2. from django.contrib.messages.storage.base import BaseStorage
  3. from django.contrib.messages.storage.cookie import MessageDecoder, MessageEncoder
  4. from django.core.exceptions import ImproperlyConfigured
  5. class SessionStorage(BaseStorage):
  6. """
  7. Store messages in the session (that is, django.contrib.sessions).
  8. """
  9. session_key = "_messages"
  10. def __init__(self, request, *args, **kwargs):
  11. if not hasattr(request, "session"):
  12. raise ImproperlyConfigured(
  13. "The session-based temporary message storage requires session "
  14. "middleware to be installed, and come before the message "
  15. "middleware in the MIDDLEWARE list."
  16. )
  17. super().__init__(request, *args, **kwargs)
  18. def _get(self, *args, **kwargs):
  19. """
  20. Retrieve a list of messages from the request's session. This storage
  21. always stores everything it is given, so return True for the
  22. all_retrieved flag.
  23. """
  24. return (
  25. self.deserialize_messages(self.request.session.get(self.session_key)),
  26. True,
  27. )
  28. def _store(self, messages, response, *args, **kwargs):
  29. """
  30. Store a list of messages to the request's session.
  31. """
  32. if messages:
  33. self.request.session[self.session_key] = self.serialize_messages(messages)
  34. else:
  35. self.request.session.pop(self.session_key, None)
  36. return []
  37. def serialize_messages(self, messages):
  38. encoder = MessageEncoder()
  39. return encoder.encode(messages)
  40. def deserialize_messages(self, data):
  41. if data and isinstance(data, str):
  42. return json.loads(data, cls=MessageDecoder)
  43. return data