dataclassを使ったYAML形式で保存／ロード可能な設定クラス

Pythonで設定オブジェクトを定義する場合にPython3.7から追加されたdataclasses.dataclassが便利。デコレータをつけるだけでコンストラクタを省略できる。各項目に型を指定すればpyrightやmypyで型推論をして事前にバグを発見できる。今回はdataclassとして振る舞う設定オブジェクトをYAMLファイルで保存／ロードできるようにしたいので実装した。

https://github.com/kzmssk/yaml_config

dataclassによる設定ファイル

例えば、このような設定オブジェクトを定義できる。

import typing
import dataclasses


@dataclasses.dataclass
class MyConfig:
    foo: int
    bar: float
    baz: typing.List[str]


config = MyConfig(foo=1, bar=2.0, baz=['3 and 4'])

print(config.foo)  # 1

普通のobjectを継承したクラスの場合はコンストラクタを書く必要があるが、dataclassの場合は必要ない。

保存／ロードできる設定ファイル

設定ファイルはPythonファイルだけではなくてテキストファイルとして保持できると便利なことがある。今回は形式としてYAMLを使うことにする。pip install pyyamlをすればjsonパッケージと似たようにYAMLファイルを使った保存とロードができるようになる。

上記のdataclassの設定オブジェクトをYAMLファイルとして保存し、なおかつロードできるようなクラスを作りたい。ついでにpathlib.Pathは設定によくある型なのでこれもYAMLにするときにちゃんと変換できるようにしたい。

ということでYAML形式で保存／ロード可能な設定クラスをdataclassを使って作ってみた：

import pathlib
import dataclasses
import yaml
import inspect

@dataclasses.dataclass
class YamlConfig:
    def save(self, config_path: pathlib.Path):
        """ Export config as YAML file """
        assert config_path.parent.exists(), f'directory {config_path.parent} does not exist'

        def convert_dict(data):
            for key, val in data.items():
                if isinstance(val, pathlib.Path):
                    data[key] = str(val)
                if isinstance(val, dict):
                    data[key] = convert_dict(val)
            return data

        with open(config_path, 'w') as f:
            yaml.dump(convert_dict(dataclasses.asdict(self)), f)

    @classmethod
    def load(cls, config_path: pathlib.Path):
        """ Load config from YAML file """
        assert config_path.exists(), f'YAML config {config_path} does not exist'

        def convert_from_dict(parent_cls, data):
            for key, val in data.items():
                child_class = parent_cls.__dataclass_fields__[key].type
                if child_class == pathlib.Path:
                    data[key] = pathlib.Path(val)
                if inspect.isclass(child_class) and issubclass(child_class, YamlConfig):
                    data[key] = child_class(**convert_from_dict(child_class, val))
            return data

        with open(config_path) as f:
            config_data = yaml.full_load(f)
            # recursively convert config item to YamlConfig
            config_data = convert_from_dict(cls, config_data)
            return cls(**config_data)

テストを含めたコードはGitHubのレポジトリにあげてある。

自分の設定オブジェクトを定義したいときはたとえば以下のようにすれば良い：

#import YamlConfig and typing
@dataclasses.dataclass
class MyConfig(YamlConfig):
    val_float: float
    val_list: typing.List[int]
    val_str: str

YamlConfigを継承すればsaveとloadを使うことができる。

# import MyConfig and pathlib
config = MyConfig(val_float=1.0, val_list=[1,2], val_str='3')
config.save(pathlib.Path('./my_config.yaml')

すでにあるYAMLファイルから設定を読み込みたい場合は

# import MyConfig and pathlib
config = MyConfig.load(pathlib.Path('./my_config.yaml'))

とすればよい。

さらに設定の子設定を作るなどの階層構造を定義した場合も saveとloadが動く。

@dataclasses.dataclass
class MySubSubConfig(YamlConfig):
    val_float: float
    val_list: typing.List[int]
    val_str: str


@dataclasses.dataclass
class MySubConfig(YamlConfig):
    val_int: int
    val_path: pathlib.Path
    sub_sub_config: MySubSubConfig


@dataclasses.dataclass
class MyConfig(YamlConfig):
    sub_config: MySubConfig

このようなMyConfigのloadとsaveもできる。

感想など

アトミックな型とpathlib.Pathは対応しているが他の型（自作クラスなど）は使えない。またsaveとloadの実装はyaml.add_representerやyaml.add_constructorを使った方がよかったのかもしれない。