PythonのDataclassとYAMLの組み合わせ

PythonのDataclassとYAMLを組み合わせることで、設定オブジェクトを簡単に定義し、保存やロードを行うことができます。

まずは、Python3.7から追加されたdataclasses.dataclassを使って設定オブジェクトを定義します。このデコレータをつけるだけでコンストラクタを省略でき、各項目に型を指定すればpyrightやmypyで型推論をして事前にバグを発見できます。

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

次に、この設定オブジェクトをYAMLファイルとして保存し、ロードできるようにします。これにはpyyamlパッケージを使用します。

import pathlib
import dataclasses
import yaml
import inspect

@dataclasses.dataclass
class YamlConfig:
    def save(self, config_path: pathlib.Path):
        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):
        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)
            config_data = convert_from_dict(cls, config_data)
            return cls(**config_data)

以上のようにすると、YamlConfigを継承した設定オブジェクトをYAMLファイルとして保存し、ロードすることができます。

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

既存のYAMLファイルから設定を読み込む場合は、以下のようにします。

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

このようにPythonのDataclassとYAMLを組み合わせることで、設定オブジェクトの定義、保存、ロードを簡単に行うことができます。設定の子設定を作るなどの階層構造を定義した場合もsaveとloadが動作します。