enumとdataclassesを含んだ値をテキトーにJSONとしてseiralize/deserializeしたい

昔に似たようなタイトルの記事を書いていましたが、これとはちょっと違った内容です。

datacalssesのインスタンスを含んだ値をJSONにするときのメモ

enumやdataclassesを含んだ値をテキトーにJSONとしてserialize/desserializeしたくなった。パフォーマンスは全く気にしなくて良い。どちらかといえば、schemaの定義なしにJSON側にpythonのオブジェクトの構造に関する情報を持たせたいというニュアンスのほうが強い。テキストファイルとしてエディタで開いて出力された値を書き換えたりなどしたかったので、pickleではダメだった。

とりあえず、以下の様なオブジェクトが一致することを目指す。

import enum
import dataclasses


class Color(enum.Enum):
    Red = 1
    Green = 2
    Blue = 3
    Yellow = 4


@dataclasses.dataclass
class Person:
    name: str
    initial: str = dataclasses.field(init=False)
    color: Color

    def __post_init__(self):
        self.initial = self.name.upper()[0]


p = Person(name="foo", color=Color.Green)
print(p == loads(dumps(p)))

雑にまとめるとこんな感じ。

dumps(),loads()を定義したい
init=Falseなフィールドが存在する
enumのフィールドが存在する
dataclassesを使ったフィールドが存在する（上の例ではフィールドにはdataclassesが使われていないが、実際には使われている）

これに関するserialize/deserializeをでっちあげたかった。

dumps

enumは __enum__ というフィールドでwrapして出力することにする。ペアとなる値は定義したEnumの属性名。例えば、Color.Greenは以下の様に出力される。

{"__enum__": "__main__.Color.Green"}

dataclasssesは __dataclass__ というフィールドでwrapして出力することにする。こちらはどのようなデータかという情報も欲しいので、typeとvalueというフィールドを持った辞書にする。どのモジュールに所属しているかの情報もほしいのでフルパスでtypeは出力する

{
  "__dataclass__": {
    "type": "__main__.Person",
    "value": {
      "name": "foo",
      "initial": "F",
      "color": {
        "__enum__": "__main__.Color.Green"
      }
    }
  }
}

init=Falseなフィールドは、pythonのコンストラクターを呼び出したときに利用できない。自動で生成される __init__() の引数として扱えない。 dataclasses.fields()を利用して、対象のフィールドを探してpopする。

実装はこんな感じ。手抜きの実装をするときに、キーワード引数を初期値部分にキャッシュ用のdictを入れるのが便利（二度とdeleteはできなくなるけど）。

import enum
import json
import dataclasses


def dumps(ob):
    return json.dumps(ob, indent=2, default=_default)


def _default(ob, *, _cache={}):
    if dataclasses.is_dataclass(ob):
        d = dataclasses.asdict(ob)
        omit_keys = _cache.get(ob.__class__)
        if omit_keys is None:
            omit_keys = _cache[ob.__class__] = [
                f.name for f in dataclasses.fields(ob) if not f.init
            ]
        for k in omit_keys:
            d.pop(k)
        return {
            "__dataclass__": {
                "type": f"{ob.__class__.__module__}.{ob.__class__.__name__}",
                "value": d,
            }
        }
    elif isinstance(ob, enum.Enum):
        return {
            "__enum__": f"{ob.__class__.__module__}.{ob.__class__.__name__}.{ob.name}",
        }
    raise TypeError(f"unexpected type {ob!r}")

loads

dumpsを作ったタイミングでwrapした__enum__と__dataclass__に対する特別扱いを入れる。JSONのオブジェクト部分のhookとしてobject_pairs_hookが用意されている。ここにdictではなく自作の関数を定義する。あと特筆する事があるとしたら、sys.modulesとimportlib.import_moduleを見て、モジュールをimportすることくらい。

実装は以下の様な感じ。

import sys
import json
from importlib import import_module


def loads(s):
    return json.loads(s, object_pairs_hook=_on_pairs)


def _on_pairs(itr):
    d = {k: v for k, v in itr}
    if "__enum__" in d:
        module, clsname, attr = d["__enum__"].rsplit(".", 3)
        m = sys.modules.get(module)
        if m is None:
            m = import_module(module)
        cls = getattr(m, clsname)
        return getattr(cls, attr)
    elif "__dataclass__" in d:
        v = d["__dataclass__"]
        module, clsname = v["type"].rsplit(".", 2)
        m = sys.modules.get(module)
        if m is None:
            m = import_module(module)
        cls = getattr(m, clsname)
        return cls(**v["value"])
    else:
        return d

実行

実際動く。validationなどは無い。

if __name__ == "__main__":

    class Color(enum.Enum):
        Red = 1
        Green = 2
        Blue = 3
        Yellow = 4

    @dataclasses.dataclass
    class Person:
        name: str
        initial: str = dataclasses.field(init=False)
        color: Color

        def __post_init__(self):
            self.initial = self.name.upper()[0]

    p = Person(name="foo", color=Color.Green)
    print(dumps(p))
    s = """
{
  "__dataclass__": {
    "type": "__main__.Person",
    "value": {
      "name": "foo",
      "color": {
        "__enum__": "__main__.Color.Green"
      }
    }
  }
}
"""
    print(loads(s))
    print(p == loads(dumps(p)))

実行結果

{
  "__dataclass__": {
    "type": "__main__.Person",
    "value": {
      "name": "foo",
      "color": {
        "__enum__": "__main__.Color.Green"
      }
    }
  }
}
Person(name='foo', initial='F', color=<Color.Green: 2>)
True

細々とした話

init=Falseにするくらいならpropertyを使ってはどうか？ -> 値を更新したくなる。setterを書こうとすると保存する属性をアレコレするのが面倒
本当は更新した値をserializeしたあとにその値を使いたくない？ -> 使いたい（頑張ればできるが。。）
結局、init=Falseとか使ってしまうと対称性が壊れる。全部引数として受け入れるようにして、別途クラスメソッドなどを定義したら？　-> めんどくさい
popダサくない？ -> ダサい

(もちろん、frozen=Trueにしてdataclasses.replaceを利用というimmutableっぽい使い方ができるならやっている。何も無いところからはじめるならそれも検討に入れる)