diff --git a/lemoncurry/msgpack.py b/lemoncurry/msgpack.py index 6328f2c..4a5fb57 100644 --- a/lemoncurry/msgpack.py +++ b/lemoncurry/msgpack.py @@ -1,11 +1,24 @@ import msgpack +import pickle from django_redis.serializers.base import BaseSerializer +def default(obj): + # Pickle anything that MessagePack can't handle itself. + return msgpack.ExtType(69, pickle.dumps(obj, protocol=4)) + + +def ext_hook(code, data): + # Unpickle if we pickled - otherwise do nothing. + if code == 69: + return pickle.loads(data) + return msgpack.ExtType(code, data) + + class MSGPackModernSerializer(BaseSerializer): def dumps(self, value): - return msgpack.packb(value, use_bin_type=True) + return msgpack.packb(value, default=default, use_bin_type=True) def loads(self, value): - return msgpack.unpackb(value, raw=False) + return msgpack.unpackb(value, ext_hook=ext_hook, raw=False)