42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
|
import os
|
||
|
|
import re
|
||
|
|
import typing
|
||
|
|
from typing import Any, TextIO
|
||
|
|
|
||
|
|
from yaml import SafeLoader
|
||
|
|
|
||
|
|
_env_replace_matcher = re.compile(r"\$\{(\w|_)+:?.*}")
|
||
|
|
|
||
|
|
|
||
|
|
@typing.no_type_check # pyaml does not have good hints, everything is Any
|
||
|
|
def load_yaml_with_envvars(
|
||
|
|
stream: TextIO, environ: dict[str, Any] = os.environ
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Load yaml file with environment variable expansion.
|
||
|
|
|
||
|
|
The pattern ${VAR} or ${VAR:default} will be replaced with
|
||
|
|
the value of the environment variable.
|
||
|
|
"""
|
||
|
|
loader = SafeLoader(stream)
|
||
|
|
|
||
|
|
def load_env_var(_, node) -> str:
|
||
|
|
"""Extract the matched value, expand env variable, and replace the match."""
|
||
|
|
value = str(node.value).removeprefix("${").removesuffix("}")
|
||
|
|
split = value.split(":", 1)
|
||
|
|
env_var = split[0]
|
||
|
|
value = environ.get(env_var)
|
||
|
|
default = None if len(split) == 1 else split[1]
|
||
|
|
if value is None and default is None:
|
||
|
|
raise ValueError(
|
||
|
|
f"Environment variable {env_var} is not set and not default was provided"
|
||
|
|
)
|
||
|
|
return value or default
|
||
|
|
|
||
|
|
loader.add_implicit_resolver("env_var_replacer", _env_replace_matcher, None)
|
||
|
|
loader.add_constructor("env_var_replacer", load_env_var)
|
||
|
|
|
||
|
|
try:
|
||
|
|
return loader.get_single_data()
|
||
|
|
finally:
|
||
|
|
loader.dispose()
|