PythonでDataclassをYAMLにダンプする方法

PythonのdataclassをYAML形式で保存／ロードする方法について説明します。Python3.7から追加されたdataclasses.dataclassは、設定オブジェクトを定義する際に非常に便利です。

まず、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

次に、この設定オブジェクトをYAML形式で保存し、ロードできるようにします。そのためには、まずpyyamlパッケージをインストールします。

pip install pyyaml

そして、以下のようにYamlConfigクラスを作成します。

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)
            config_data = convert_from_dict(cls, config_data)
            return cls(**config_data)

これで、MyConfigクラスが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'))

とすればよいです。

以上がPythonのdataclassをYAML形式で保存／ロードする方法です。この方法を使えば、設定の管理が容易になります。ぜひ試してみてください。