PythonでDecimalやdatetimeをシリアライズする

Pythonで少し複雑なデータをシリアライズする時に便利な関数を作成しました．MessagePackとJSONに対応しており，srslyというシリアライズのライブラリを使用しています．

import datetime
import decimal
import json

from srsly import msgpack

#
# MessagePack
#


def msgpack_dumps(d) -> bytes:
    return msgpack.packb(d, default=_default_func, use_bin_type=True)


def _default_func(obj):
    if isinstance(obj, decimal.Decimal):
        return msgpack.ExtType(42, str(obj).encode(encoding="utf-8"))
    if isinstance(obj, datetime.datetime):
        return msgpack.ExtType(43, obj.isoformat().encode(encoding="utf-8"))
    if isinstance(obj, datetime.date):
        return msgpack.ExtType(44, obj.isoformat().encode(encoding="utf-8"))
    if isinstance(obj, datetime.time):
        return msgpack.ExtType(45, obj.isoformat().encode(encoding="utf-8"))
    raise TypeError(f"Unknown type: {obj}")


def msgpack_loads(d: bytes):
    return msgpack.unpackb(d, ext_hook=_ext_hook, raw=False)


def _ext_hook(code, d: bytes):
    if code == 42:
        return decimal.Decimal(d.decode(encoding="utf-8"))
    if code == 43:
        return datetime.datetime.fromisoformat(d.decode(encoding="utf-8"))
    if code == 44:
        return datetime.date.fromisoformat(d.decode(encoding="utf-8"))
    if code == 45:
        return datetime.time.fromisoformat(d.decode(encoding="utf-8"))
    raise TypeError(f"Unknown type: {code}, {d}")


#
# JSON encoder/decoder
#


class JSONEncoder(json.JSONEncoder):

    def default(self, obj):
        if isinstance(obj, decimal.Decimal):
            return {
                "__type__": "Decimal",
                "value": str(obj),
            }
        if isinstance(obj, datetime.datetime):
            return {
                "__type__": "datetime",
                "value": [
                    obj.year,
                    obj.month,
                    obj.day,
                    obj.hour,
                    obj.minute,
                    obj.second,
                ],
            }
        if isinstance(obj, datetime.date):
            return {
                "__type__": "date",
                "value": [obj.year, obj.month, obj.day],
            }
        if isinstance(obj, datetime.time):
            return {
                "__type__": "time",
                "value": [obj.hour, obj.minute, obj.second],
            }
        return super().default(obj)


class JSONDecoder(json.JSONDecoder):
    def __init__(self, *args, **kwargs):
        super().__init__(object_hook=self.object_hook, *args,  **kwargs)

    def object_hook(self, obj):
        v = obj.get("__type__")
        if v is None:
            return obj
        if v == "Decimal":
            return decimal.Decimal(obj["value"])
        if v == "datetime":
            return datetime.datetime(*obj["value"])
        if v == "date":
            return datetime.date(*obj["value"])
        if v == "time":
            return datetime.time(*obj["value"])
        raise TypeError(f"Unserializable object {obj} of type {type(obj)}")

unittest

import datetime
import decimal
import json
import unittest

from ser import msgpack_dumps, msgpack_loads, JSONDecoder, JSONEncoder

tests = [
    {},
    [1, 2, 3, 4],
    {
        "a": 1,
        "b": decimal.Decimal(2),
        "c": datetime.datetime(2000, 1, 1, 14, 15, 16),
        "d": datetime.date(2001, 2, 3),
        "e": datetime.time(11, 12, 13),
        "f": [1, 2, 3, decimal.Decimal(4)],
    },
]


class SerializeTest(unittest.TestCase):
    def test_msgpack(self):
        for o in tests:
            with self.subTest(o=o):
                b = msgpack_dumps(o)
                p = msgpack_loads(b)
                self.assertEqual(isinstance(b, bytes), True)
                self.assertEqual(p, o)

    def test_json(self):
        for o in tests:
            with self.subTest(o=o):
                b = json.dumps(o, cls=JSONEncoder)
                p = json.loads(b, cls=JSONDecoder)
                self.assertEqual(isinstance(b, str), True)
                self.assertEqual(p, o)

これを実行し，動作が確認できました．

% sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
% python3 --version
Python 3.9.7
% python -m unittest                             
..
----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK