fallback.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. from django.contrib.messages.storage.base import BaseStorage
  2. from django.contrib.messages.storage.cookie import CookieStorage
  3. from django.contrib.messages.storage.session import SessionStorage
  4. class FallbackStorage(BaseStorage):
  5. """
  6. Try to store all messages in the first backend. Store any unstored
  7. messages in each subsequent backend.
  8. """
  9. storage_classes = (CookieStorage, SessionStorage)
  10. def __init__(self, *args, **kwargs):
  11. super().__init__(*args, **kwargs)
  12. self.storages = [
  13. storage_class(*args, **kwargs) for storage_class in self.storage_classes
  14. ]
  15. self._used_storages = set()
  16. def _get(self, *args, **kwargs):
  17. """
  18. Get a single list of messages from all storage backends.
  19. """
  20. all_messages = []
  21. for storage in self.storages:
  22. messages, all_retrieved = storage._get()
  23. # If the backend hasn't been used, no more retrieval is necessary.
  24. if messages is None:
  25. break
  26. if messages:
  27. self._used_storages.add(storage)
  28. all_messages.extend(messages)
  29. # If this storage class contained all the messages, no further
  30. # retrieval is necessary
  31. if all_retrieved:
  32. break
  33. return all_messages, all_retrieved
  34. def _store(self, messages, response, *args, **kwargs):
  35. """
  36. Store the messages and return any unstored messages after trying all
  37. backends.
  38. For each storage backend, any messages not stored are passed on to the
  39. next backend.
  40. """
  41. for storage in self.storages:
  42. if messages:
  43. messages = storage._store(messages, response, remove_oldest=False)
  44. # Even if there are no more messages, continue iterating to ensure
  45. # storages which contained messages are flushed.
  46. elif storage in self._used_storages:
  47. storage._store([], response)
  48. self._used_storages.remove(storage)
  49. return messages