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形式で保存/ロードする方法です。この方法を使えば、設定の管理が容易になります。ぜひ試してみてください。