url_safe.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. import typing as _t
  2. import zlib
  3. from ._json import _CompactJSON
  4. from .encoding import base64_decode
  5. from .encoding import base64_encode
  6. from .exc import BadPayload
  7. from .serializer import Serializer
  8. from .timed import TimedSerializer
  9. class URLSafeSerializerMixin(Serializer):
  10. """Mixed in with a regular serializer it will attempt to zlib
  11. compress the string to make it shorter if necessary. It will also
  12. base64 encode the string so that it can safely be placed in a URL.
  13. """
  14. default_serializer = _CompactJSON
  15. def load_payload(
  16. self,
  17. payload: bytes,
  18. *args: _t.Any,
  19. serializer: _t.Optional[_t.Any] = None,
  20. **kwargs: _t.Any,
  21. ) -> _t.Any:
  22. decompress = False
  23. if payload.startswith(b"."):
  24. payload = payload[1:]
  25. decompress = True
  26. try:
  27. json = base64_decode(payload)
  28. except Exception as e:
  29. raise BadPayload(
  30. "Could not base64 decode the payload because of an exception",
  31. original_error=e,
  32. ) from e
  33. if decompress:
  34. try:
  35. json = zlib.decompress(json)
  36. except Exception as e:
  37. raise BadPayload(
  38. "Could not zlib decompress the payload before decoding the payload",
  39. original_error=e,
  40. ) from e
  41. return super().load_payload(json, *args, **kwargs)
  42. def dump_payload(self, obj: _t.Any) -> bytes:
  43. json = super().dump_payload(obj)
  44. is_compressed = False
  45. compressed = zlib.compress(json)
  46. if len(compressed) < (len(json) - 1):
  47. json = compressed
  48. is_compressed = True
  49. base64d = base64_encode(json)
  50. if is_compressed:
  51. base64d = b"." + base64d
  52. return base64d
  53. class URLSafeSerializer(URLSafeSerializerMixin, Serializer):
  54. """Works like :class:`.Serializer` but dumps and loads into a URL
  55. safe string consisting of the upper and lowercase character of the
  56. alphabet as well as ``'_'``, ``'-'`` and ``'.'``.
  57. """
  58. class URLSafeTimedSerializer(URLSafeSerializerMixin, TimedSerializer):
  59. """Works like :class:`.TimedSerializer` but dumps and loads into a
  60. URL safe string consisting of the upper and lowercase character of
  61. the alphabet as well as ``'_'``, ``'-'`` and ``'.'``.
  62. """