fastapi-cache/fastapi_cache/coder.py

import codecs
import datetime
import json
import pickle  # nosec:B403
from decimal import Decimal
from typing import Any, Callable, ClassVar, Dict, TypeVar, overload

import pendulum
from fastapi.encoders import jsonable_encoder
from pydantic import BaseConfig, ValidationError, fields
from starlette.responses import JSONResponse
from starlette.templating import _TemplateResponse as TemplateResponse

_T = TypeVar("_T")


CONVERTERS: dict[str, Callable[[str], Any]] = {
    "date": lambda x: pendulum.parse(x, exact=True),
    "datetime": lambda x: pendulum.parse(x, exact=True),
    "decimal": Decimal,
}


class JsonEncoder(json.JSONEncoder):
    def default(self, obj: Any) -> Any:
        if isinstance(obj, datetime.datetime):
            return {"val": str(obj), "_spec_type": "datetime"}
        elif isinstance(obj, datetime.date):
            return {"val": str(obj), "_spec_type": "date"}
        elif isinstance(obj, Decimal):
            return {"val": str(obj), "_spec_type": "decimal"}
        else:
            return jsonable_encoder(obj)


def object_hook(obj: Any) -> Any:
    _spec_type = obj.get("_spec_type")
    if not _spec_type:
        return obj

    if _spec_type in CONVERTERS:
        return CONVERTERS[_spec_type](obj["val"])
    else:
        raise TypeError("Unknown {}".format(_spec_type))


class Coder:
    @classmethod
    def encode(cls, value: Any) -> str:
        raise NotImplementedError

    @classmethod
    def decode(cls, value: str) -> Any:
        raise NotImplementedError

    # (Shared) cache for endpoint return types to Pydantic model fields.
    # Note that subclasses share this cache! If a subclass overrides the
    # decode_as_type method and then stores a different kind of field for a
    # given type, do make sure that the subclass provides its own class
    # attribute for this cache.
    _type_field_cache: ClassVar[Dict[Any, fields.ModelField]] = {}

    @overload
    @classmethod
    def decode_as_type(cls, value: str, type_: _T) -> _T:
        ...

    @overload
    @classmethod
    def decode_as_type(cls, value: str, *, type_: None) -> Any:
        ...

    @classmethod
    def decode_as_type(cls, value: str, *, type_: _T | None) -> _T | Any:
        """Decode value to the specific given type

        The default implementation uses the Pydantic model system to convert the value.

        """
        result = cls.decode(value)
        if type_ is not None:
            try:
                field = cls._type_field_cache[type_]
            except KeyError:
                field = cls._type_field_cache[type_] = fields.ModelField(
                    name="body", type_=type_, class_validators=None, model_config=BaseConfig
                )
            result, errors = field.validate(result, {}, loc=())
            if errors is not None:
                if not isinstance(errors, list):
                    errors = [errors]
                raise ValidationError(errors, type_)
        return result


class JsonCoder(Coder):
    @classmethod
    def encode(cls, value: Any) -> str:
        if isinstance(value, JSONResponse):
            return value.body.decode()
        return json.dumps(value, cls=JsonEncoder)

    @classmethod
    def decode(cls, value: str) -> str:
        return json.loads(value, object_hook=object_hook)


class PickleCoder(Coder):
    @classmethod
    def encode(cls, value: Any) -> str:
        if isinstance(value, TemplateResponse):
            value = value.body
        return codecs.encode(pickle.dumps(value), "base64").decode()

    @classmethod
    def decode(cls, value: str) -> Any:
        return pickle.loads(codecs.decode(value.encode(), "base64"))  # nosec:B403,B301

    @classmethod
    def decode_as_type(cls, value: str, *, type_: Any) -> Any:
        # Pickle already produces the correct type on decoding, no point
        # in paying an extra performance penalty for pydantic to discover
        # the same.
        return cls.decode(value)