diff --git a/fastapi_cache/coder.py b/fastapi_cache/coder.py index b6cd683..15ed118 100644 --- a/fastapi_cache/coder.py +++ b/fastapi_cache/coder.py @@ -1,8 +1,9 @@ +import codecs import datetime import json import pickle # nosec:B403 from decimal import Decimal -from typing import Any, Dict, Union +from typing import Any import pendulum from fastapi.encoders import jsonable_encoder @@ -44,7 +45,7 @@ class Coder: raise NotImplementedError @classmethod - def decode(cls, value: Any) -> Any: + def decode(cls, value: str) -> Any: raise NotImplementedError @@ -54,7 +55,7 @@ class JsonCoder(Coder): return json.dumps(value, cls=JsonEncoder) @classmethod - def decode(cls, value: Any) -> str: + def decode(cls, value: str) -> str: return json.loads(value, object_hook=object_hook) @@ -63,8 +64,8 @@ class PickleCoder(Coder): def encode(cls, value: Any) -> str: if isinstance(value, TemplateResponse): value = value.body - return str(pickle.dumps(value)) + return codecs.encode(pickle.dumps(value), "base64").decode() @classmethod - def decode(cls, value: Any) -> Any: - return pickle.loads(bytes(value)) # nosec:B403,B301 + def decode(cls, value: str) -> Any: + return pickle.loads(codecs.decode(value.encode(), "base64")) # nosec:B403,B301 diff --git a/tests/test_codecs.py b/tests/test_codecs.py new file mode 100644 index 0000000..371172d --- /dev/null +++ b/tests/test_codecs.py @@ -0,0 +1,22 @@ +from typing import Any + +import pytest + +from fastapi_cache.coder import PickleCoder + + +@pytest.mark.parametrize( + "value", + [ + 1, + "some_string", + (1, 2), + [1, 2, 3], + {"some_key": 1, "other_key": 2}, + ], +) +def test_pickle_coder(value: Any) -> None: + encoded_value = PickleCoder.encode(value) + assert isinstance(encoded_value, str) + decoded_value = PickleCoder.decode(encoded_value) + assert decoded_value == value