dataclassとobject_hookでjsonとして保存しやすいデータクラスを作る

はじめに

• dataclassとjsonを相互変換できる仕組みを自作したときの話。
  • 前もmatchを使ってやったことがあるが、あれだと定義を作るのがかなり面倒だったので作り直したもの。
• 使ったことはないがdataclass-jsonという便利そうなものがあるので、それが使えるならそっちでもいいと思う。

基本的な話

• json -> classはobject_hookで行う。
• class -> jsonはdefaultで行う。
• jsonに変換可能なクラスを作り、それを継承させることで相互変換を実現させる。
  • 継承させた先でdataclassのデコレータを付ける。
  • enumの場合はEnumより先にJsonableを継承させる必要がある。

json変換可能なクラスを定義する

• jsonと相互変換できるクラスを登録するのに__init_subclass__を利用する。
  • これでこのクラスを継承したクラスを作るたびにクラス名をキーとしてクラスが登録されるので、jsonからクラスに戻すときにここで登録したクラスを利用することができる。
  • 別モジュールに同名のクラスがあるとおかしくなる気もするが、同じ名前のクラスをjsonに変換できるクラスとして複数作ることはないと思うので多分大丈夫。
    • 心配なら__module__とかを繋げて名前をかぶらないようにする(フォルダ構成を変えると復元できなくなるデメリット付き)。

class Jsonable:
    __children__: dict[str, Any] = {}

    def __init_subclass__(cls) -> None:
        cls.__children__[dataclass.__name__] = dataclass

dataclass -> jsonを考える

• dataclassをjsonにするには、dataclasses.fieldsを使えば簡単にできる。
• 復元用に元のクラス名を保存したいので名前がかぶらないように"__"で囲んで基本形とする。
• その後、dataclassかEnumかを判断して個別の項目を足して返す。

def as_json(__value: Any, /) -> Any:
    if isinstance(__value, Jsonable):
        base: dict[str, Any] = {"__type__": type(__value).__name__}
        if dataclasses.is_dataclass(__value):
            return base | {field.name: getattr(__value, field.name) for field in dataclasses.fields(__value)}
        elif isinstance(__value, Enum):
            return base | {"__name__": __value.name}
        else:
            raise TypeError(f"not dataclass or enum: {type(__value).__name__!r}")
    raise TypeError(f"unknown type: {type(__value).__name__!r}")

json -> dataclassを考える

• 今度はjsonを元にdataclassへの復元方法を考える。
• text: strなフィールドを持つFooクラスの場合は以下のようなjsonに変換されてる。

  {"__type__": "Foo", "text": "sample_text"}
  

• これの復元は逆にFooに対して"type"を取り除いてキーワード引数として与えれば行うことができる。
• こちらはobject_hookなのでdict[str, Any]を受け取ることが決まっているので、今回の形式の目印である"type"を持ったオブジェクトを対象に処理を行う。
• サブクラス作成時に登録した辞書からクラスを取り出し、dataclassかEnumかを判定していい感じに処理する。
  • dataclassの場合は"_"で始まるフィールドは特殊な項目なので、キーワード引数として与えないように除外してインスタンスを作る。
  • enumの場合は名前も必要なのでそちらの所持も確認し、getattrで取得を行う。

def from_json(__value: dict[str, Any], /) -> Any:
    if typename := __value.get("__type__", None):
        cls = Jsonable.__children__[typename]
        if dataclasses.is_dataclass(cls):
            return cls(**{k: v for k, v in __value.items() if not str(k).startswith("_")})
        elif issubclass(cls, Enum) and "__name__" in __value:
            return getattr(cls, __value["__name__"])
        raise TypeError(f"unknown type: {cls.__name__!r}")
    else:
        return {k: v for k, v in __value.items()}

オマケ機能

• dataclassにするときにslotsを指定することでdatetime.dateなども継承できるようになる。
  • 日付や日時もこちらを継承した日付型などにすれば個別にdefault用の処理を書かないで済む。
  • ちゃんとDate.today()もできて、継承もしているので型チェックも通り、普通のdateと同じように扱える。

  @dataclasses.dataclass(frozen=True, slots=True)
  class Date(Jsonable, date):
      year: int
      month: int
      day: int
  

  • ※dateと同じようにイミュータブルにしておく必要があるのでfrozenもセットで指定する。

拡張

• Jsonableクラスにsaveメソッドを持たせておいたりすると、継承した小さいクラスをとりあえずファイルに書き出して確認したい！みたいな簡単な動作確認のときに毎回jsonやらなにやらをインポートする手間が省ける。
  • 必要に応じてensure_asciiとかindentもいい感じに指定しておくと扱いが楽。
• 設定ファイルのルートで使う分にも便利かもしれない。
• サンプルではファイル名を省略した場合は「クラス名.json」で保存するようにしている。

class Jsonable:
    def save(self, __path: str | None = None, /, *, default: Callable[[Any], Any]=as_object, **kwargs: Any) -> None:
        path = __path or f"{type(self).__name__}.json"
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self, f, default=default, **kwargs)

コード全体

from __future__ import annotations

import dataclasses
import json
from datetime import date
from enum import Enum, auto
from typing import Any


class Jsonable:
    __children__: dict[str, Any] = {}

    def __init_subclass__(cls) -> None:
        cls.__children__[cls.__name__] = cls


def as_json(__value: Any, /) -> Any:
    if isinstance(__value, Jsonable):
        base: dict[str, Any] = {"__type__": type(__value).__name__}
        if dataclasses.is_dataclass(__value):
            return base | {field.name: getattr(__value, field.name) for field in dataclasses.fields(__value)}
        elif isinstance(__value, Enum):
            return base | {"__name__": __value.name}
        else:
            raise TypeError(f"not dataclass or enum: {type(__value).__name__!r}")
    raise TypeError(f"unknown type: {type(__value).__name__!r}")


def from_json(__value: dict[str, Any], /) -> Any:
    if typename := __value.get("__type__", None):
        cls = Jsonable.__children__[typename]
        if dataclasses.is_dataclass(cls):
            return cls(**{k: v for k, v in __value.items() if not str(k).startswith("_")})
        elif issubclass(cls, Enum) and "__name__" in __value:
            return getattr(cls, __value["__name__"])
        raise TypeError(f"unknown type: {cls.__name__!r}")
    else:
        return {k: v for k, v in __value.items()}


@dataclasses.dataclass
class Hoge(Jsonable):
    foo: str
    fugas: list[Fuga]
    fugamap: dict[str, Fuga]
    x: Samples


@dataclasses.dataclass
class Fuga(Jsonable):
    name: str
    qty: int


@dataclasses.dataclass
class Foo(Jsonable):
    text: str


@dataclasses.dataclass(frozen=True, slots=True)
class Date(Jsonable, date):
    year: int
    month: int
    day: int


class Samples(Jsonable, Enum):
    MEM_1 = auto()
    MEM_2 = auto()


def main():
    obj = [Hoge("hoge", [Fuga("fuga", 5)], {"x": Fuga("X", 3)}, Samples.MEM_2), Foo("test"), Date.today()]
    path = "sample.json"
    with open(path, "w", encoding="utf-8") as w:
        json.dump(obj, w, ensure_ascii=False, indent=4, default=as_json)
    with open(path, encoding="utf-8") as r:
        new_obj = json.load(r, object_hook=from_json)
    print(obj == new_obj, new_obj)


if __name__ == "__main__":
    main()