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