PythonでYAMLの内容を自動補完する

この記事の概要

YAMLファイルに何らかの設定値を書き、Pythonからそれをロードして参照することがあると思います。特に機械学習分野だと多い印象があります。

しかし、ロードした結果は辞書オブジェクトとなっており、そのキー名や、階層構造はYAMLファイルを目で見て確認する必要があります。
私はコーディング中にYAML確認作業が入るのが我慢できませんでした。

この記事は、YAMLファイルでconfigを定義すれば、Python実装中もその項目が自動補完されてほしい人向けにその実現方法の1つを紹介しています。

コードはここに置いてあります。このコードはモジュール化されており、pipでインストール可能です。

YAMLファイルに設定値を記述する

まず、YAMLファイルに設定値を記述する例について見ていきます。

config.yml

model:
    in_channels: 3
    n_blocks: 10
    block:
        channels: 64
        activation: relu
    out_channels: 1

YAMLファイルには上のように階層的に設定値を記述する事が多いかなと思います。
これをPythonスクリプトから読み込むとき、多少書き方は違うかもしれませんが、だいたいこんな感じになると思います。

main.py

config = yaml.safe_load('./config.yml')
model = Model(
    in_channels = config['model']['in_channels'],
    out_channels = config['model']['out_channels'],
    n_blocks = config['model']['n_blobks'],
    block_channels = config['model']['block']['channels'],
    ...
)

まず第一に自分はこの辞書のキー名をハードコーディングするという作業が耐えられませんでした。第二に、YAMLで定義した階層構造について、どの値がどの階層で定義されているかを思い出すのにYAMLファイルを見に行くのが耐えられませんでした。

理想的なのは、YAMLファイルでconfigを定義すれば、Pythonスクリプトで実装中もYAMLで設定した項目が自動補完されるというのが望ましいです。

main.py[理想]

config = get_config('./config.yml')
model = Model(
    in_channels = config.model.in_channels,
    out_channels = config.model.out_channels, # <- ドットアクセスすると自動補完される。なんなら型も確認してほしい

YACSではYAMLファイルの雛形をPythonスクリプトで定義するため、YAMLの内容をPythonスクリプト内で自動補完で出させるのは可能なんじゃないかと思いましたが、自分のPyCharmでは補完されませんでした。

yacs

from yacs.config import CfgNode as CN

_C = CN()

_C.MODEL = CN()
_C.MODEL.IN_CHANNELS = 3

def get_cfg_defaults():
  return _C.clone()

config = get_cfg_defaults()
config.MODEL.IN_CHANNELS # <-- 自動補完されない

なので、YAMLファイルでconfigを定義すれば、Pythonスクリプトで実装中もYAMLで設定した項目が自動補完されるフレームワークを自作することにしました。

YAMLファイルから対応するPythonスクリプトを生成する

まず、YAMLファルから同じ内容のPythonファイルを自動生成できれば、実装する際はPythonファイルを参照することで自動補完が可能になると考えました。こんなイメージです。

config.yml

hoge: piyo

YAML -> Python 変換

config.py

hoge: str = 'piyo'

しかし上のような単純な変換では、YAMLでの階層構造をPython上に表現することが難しいと感じました。

そこで、もう少し複雑にして、階層構造は全てデータクラス化することで対応することにしました。

config.yml

model:
    in_channels: 3
    n_blocks: 10
    block:
        channels: 64
        activation: relu
    out_channels: 1

YAML -> Python 変換

config.py

@dataclass
class ModelBlock:
    channels: int = 64
    activation: str = 'relu'

@dataclass
class Model:
    in_channels: int = 3
    n_blocks: int = 10
    block: ModelBlock = ModelBlock()
    out_channels: int = 1

@dataclass
class Config:
    model: Model = Model()

Config().model.block.channels  # <- 全て補完可能

データクラスにした理由は、マジックメソッドの自動生成によって記述量が少なくなることと、frozen=Trueにすることでメンバ変数をうっかり変更する事故を防ぐためです。
ついでに型アノテーションも生成するようにしています。ここらへんは再帰で実装しています。
こんな感じです。

これだけだとあくまでYAMLファイルをPythonスクリプトに変換するだけですが、せっかくconfigをクラス化できそたので、設定にまつわる便利メソッドもいくつか付与した上で生成したいです。

文字列結合だけで生成を書いていくのはつらそうだったので、サードパーティ生のPythonスクリプト生成モジュール、prestringを使用することにしました。

こんな感じで実装しました。
いくつかのファイルを結合しながら一つのconfig.pyを生成
するイメージです。

config.yml

model:
    in_channels: 3
    n_blocks: 10
    block:
        channels: 64
        activation: relu
    out_channels: 1

YAML -> Python 変換

config.py

@dataclass
class ModelBlock:
    channels: int = 64
    activation: str = 'relu'

@dataclass
class Model:
    in_channels: int = 3
    n_blocks: int = 10
    block: ModelBlock = ModelBlock()
    out_channels: int = 1

@dataclass
class Config:
    model: Model = Model()

    def some_cool_method():
        ...

class ConfigGenerator:
    def generate():
        ...
        return Config()

ConfigクラスがYAMLの内容を保持する責務を負っています。
ConfigGeneratorクラスを作った理由は、実際にYAMLファイルから設定値を読み込むにあたって、現在のConfigクラスと矛盾がないか、型が違わないか等をチェックする役割が必要だったからです。

これにより、目標であったYAMLファイルでconfigを定義すれば、（YAMLファイルの内容をPythonクラスに変換することで）Pythonスクリプトで実装中もYAMLで設定した項目が自動補完される事が可能になりました。

main.py

config = ConfigGenerator().generate()
model = Model(
    in_channels = config.model.in_channels,
    out_channels = config.model.out_channels, # <- ドットアクセスすると自動補完されるし型アノテーションも見てくれる

ただし、一手間として、YAMLファイルを書いたらpythonスクリプトに変換するコマンドをターミナルから叩かないといけないという制約が加わりました。

1秒の手間でtypoやYAMLファイルを見に行く作業がなくなるのなら自分は許容範囲でしたのでOKです。

追加機能

configをデータクラス化したので、ついでにいろんなクラスメソッドも自動生成で追加しておけば便利なんじゃないかなと思って追加してみました。

設定値の上書き機能

実験で頻繁に設定値を変更するようになると、「デフォルト値の設定はdefault.ymlに記述しておいて、一部の値はexp1.ymlで更新する」というような場面が出てきます。
また、視認性を高めるために「デフォルト値の設定はdefault.ymlに書いておいて、モデルの設定をmodel.ymlで、データセットの設定をdataset.yml上書きする」という場面もあります。

そういう時に

main.py

config = ConfigGenerator() \
    .update_by(['exp1.yml']) \
    .generate()
config = ConfigGenerator() \
    .update_by(['model.yml', 'dataset.yml']) \
    .generate()

と出来たら便利なのでそのような機能をConfigGeneratorに追加することにしました。
とは言え、update_byで渡されたYAMLについて、辞書としてロードしつつ、型や変数名を確認して対応する値を上書き変更するだけです。

一応、複数のYAMLファイルを並列で受け取れるようにしましたが、もしお互いのファイルで同じ設定値を異なる値に上書きしようとした場合はエラーを出すようにしました。

pprint機能

スクリプト実行時に設定値を見やすく表示することで、思わぬ事故を防ぐことが出来ます。

こんな感じで表示されるようにしました。

config = ConfigGenerator() \
    .update_by(['model.yml']) \
    .update_by(['exp01.yml']) \
    .generate()

config.pprint(wait_yes=True)  # <- 表示を確認しYESを押さない限り、コードが実行されないようにする

出力結果

default from /config/default.yml
model:
    in_channels: 3
    n_blocks: 20 (default 10, changed by /config/model.yml)
    block:
        channels: 32 (default 64, changed by /config/exp01.yml)
        activation: relu
    out_channels: 1

何らかのファイルでデフォルト値からアップデートした場合は警告メッセージを出すようにしてみました。これにより予定していた設定値と異なる設定値でコードを実行した時にすぐに気付くことが出来ます。