Pythonで設定ファイルの処理にOmegaConfを使う

今回も小ネタです。

はじめに

これまでPythonで設定ファイルっぽいものを扱うときにはyamlファイルを用意してそれを普通に読み込んで使ってました。

例えば以下のようなyamlをconfig.yamlとして用意した場合、

name: Scott
age: 18
favorite_foods: 
  - Ramen
  - Sushi
address:
  country: US
  state: CA

こんな感じで

import yaml

with open('config.yaml') as file:
    config = yaml.safe_load(file.read())

print("name: " + config['name'])
print("favorite_food: " + config['favorite_foods'][0])

実行すると、

$ python sample.py 
name: Scott
favorite_food: Ramen

という感じ。これはこれでいいのだけれど、OmegaConfというものを知ったのでそれを使ってみようかと。

OmegaConfとは

github.com

OmegaConfというのはいわゆるコンフィギュレーションファイル操作のためのライブラリで、機械学習とかのハイパーパラメータ管理で使われるhydraの足回りとして使われているとのこと。

普通のyaml形式のコンフィギュレーションファイルを読んだり書いたりするのはもちろん、コンフィギュレーションファイルの構造が定義できたり、バリデーションができたりするらしい。あとは複数ソースからのマージなんかも。

あとは値へのアクセスも上記のような辞書型形式だけでなくconfig.nameみたいな形式でも使用可能。個人的にはこっちのほうが好みなので嬉しい。

そしてInterpolationもサポートしていて先に定義した値を別の場所で参照するなんてことが可能になってる。これは公式のサンプルをコピペしたものだがserver.portの値をclient.server_portに指定できたりするのだ。これは便利そう。

server:
  host: localhost
  port: 80

client:
  url: http://${server.host}:${server.port}/
  server_port: ${server.port}

これについてはカスタムな定義も可能とのこと。

やってみる

というわけで早速さっきのconfig.yamlをOmegaConfを使って読み込んでみます。事前にpip install omegaconfしておいてください。

from omegaconf import OmegaConf

config = OmegaConf.load('config.yaml')
print("name: " + config.name)
print("favorite_food: " + config.favorite_foods[0])

行数こそさほど違いはないですが見た目にシンプルになりますね。

あとは階層構造が増えても、

print("country: " + config.address.country)

みたいな感じでアクセスできるようになってます。

辞書型やリストで定義したものから読み込んで扱うことも可能ですし、それらを最終的にyamlファイルとして書き出すことも可能。コマンドラインの引数として扱うことも可能だったりします。

from omegaconf import OmegaConf

conf = OmegaConf.create({"foo": 10, "bar": 20})
OmegaConf.save(conf, 'output.yaml')

これを実行するとこんな内容のファイルが指定したパスに作成されます。

$ cat output.yaml 
foo: 10
bar: 20

公式のサンプルみたいに一時ファイルに吐き出した上で保存するのもいいかと。

設定ファイルのマージはこんな感じです。

first_config = OmegaConf.load('conf1.yaml')
second_config = OmegaConf.load('conf2.yaml')

config = OmegaConf.merge(first_config, second_config)
print(OmegaConf.to_yaml(config))

簡単ですね。Webサーバの設定ファイルを分割したいときとかそういう感じのときに使えると思います。