この記事はPython3 Advent Calendar 2020の20日目の記事です。

はじめに

コンテナなどを使っているとYAMLファイルに環境変数を埋め込んでロード時に展開したくなることがあります。PyYAMLで環境変数の展開を行なうには以下のようなコードで実現できます。

import os
import re

import yaml

ENV_PATTERN = re.compile(r'\$\{(.*)\}')
ENV_TAG = '!env_var'

yaml.add_implicit_resolver(ENV_TAG, ENV_PATTERN, None, yaml.SafeLoader)

def env_var_constructor(loader, node):
    value = loader.construct_scalar(node)

    matched = ENV_PATTERN.match(value)
    if matched is None:
        return value
    proto = matched.group(1)

    default = None
    if len(proto.split(':')) > 1:
        env_key, default = proto.split(':')
    else:
        env_key = proto
    env_val = os.environ[env_key] if env_key in os.environ else default 

    return env_val

yaml.add_constructor(ENV_TAG, env_var_constructor, yaml.SafeLoader)

上記のコードを実行した状態でSafeLoaderを使うとロード時に環境変数を展開してくれるようになります。

example = """
a: ${EXAMPLE_A:default}
b: ${EXAMPLE_B:default}
"""
os.environ["EXAMPLE_A"]="from_env"
print(yaml.safe_load(example))

結果

{'a': 'from_env', 'b': 'default'}

環境変数を設定している場合はロードしてくれて、そうでない場合はデフォルト値が設定されています。

解説

add_implicit_resolverについて

add_implicit_resolver はYAMLの特定の値に対してタグを暗黙的に付与する設定です。

YAMLにはキーと値以外にもタグという値の型を明示的に示す文法があります。
add_implicit_resolver を使って環境変数のパターン(\$\{(.*)\})にマッチする値に暗黙的に環境変数用のタグ(!env_var)を付与する設定をSafeLoaderに設定しています。

つまり例で上げたexampleというYAMLファイルは、add_implicit_resolverによって以下のような状態になります。

a: !env_var ${EXAMPLE_A:default}
b: !env_var ${EXAMPLE_B:default}

これによって環境変数の値に対してカスタムコンストラクタを適用できる状態になります。

add_constructorについて

add_constructorについて は特定のタグのついた値に対してカスタムコンストラクタを設定する関数です。

関数env_var_constructorはカスタムコンストラクタの実装になっており、${環境変数名:デフォルト値}の文字列を正規表現マッチさせて、パースし、実際に環境変数を解決して値を返す実装になっています。

この関数を!env_varのタグに対してマッピングすることで、例で上げたYAMLファイルは、最終的に環境変数が展開され以下のような状態として値が返ります。

a: from_env
b: default

参考