codekingpro/portable-devtools

1from __future__ import annotations2 3import typing as t4import zlib5 6from ._json import _CompactJSON7from .encoding import base64_decode8from .encoding import base64_encode9from .exc import BadPayload10from .serializer import _PDataSerializer11from .serializer import Serializer12from .timed import TimedSerializer13 14 15class URLSafeSerializerMixin(Serializer[str]):16    """Mixed in with a regular serializer it will attempt to zlib17    compress the string to make it shorter if necessary. It will also18    base64 encode the string so that it can safely be placed in a URL.19    """20 21    default_serializer: _PDataSerializer[str] = _CompactJSON22 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 = False31 32        if payload.startswith(b"."):33            payload = payload[1:]34            decompress = True35 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 e43 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 e52 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 = False58        compressed = zlib.compress(json)59 60        if len(compressed) < (len(json) - 1):61            json = compressed62            is_compressed = True63 64        base64d = base64_encode(json)65 66        if is_compressed:67            base64d = b"." + base64d68 69        return base64d70 71 72class URLSafeSerializer(URLSafeSerializerMixin, Serializer[str]):73    """Works like :class:`.Serializer` but dumps and loads into a URL74    safe string consisting of the upper and lowercase character of the75    alphabet as well as ``'_'``, ``'-'`` and ``'.'``.76    """77 78 79class URLSafeTimedSerializer(URLSafeSerializerMixin, TimedSerializer[str]):80    """Works like :class:`.TimedSerializer` but dumps and loads into a81    URL safe string consisting of the upper and lowercase character of82    the alphabet as well as ``'_'``, ``'-'`` and ``'.'``.83    """84