yaml with env vars

better_yaml.py

import os
import pathlib
import re
from functools import singledispatch
from typing import Any, Dict

import yaml

GROUP_CONTAINING_VAR_NAME_RE = "var"
STRING_WITH_VARS_RE = re.compile(
    r"""
        \$                  # Match the dollar sign
        \{                  # Match the opening curly brace
        (?P<var>[^{}\.\$]+)  # Capture one or more characters that are not closing curly braces
        \}                  # Match the closing curly brace
    """,
    re.VERBOSE,
)


@singledispatch
def env_constructor(
    node,
    loader: yaml.SafeLoader,
    *,
    string_with_vars_re: re.Pattern = STRING_WITH_VARS_RE,
    group_name: str = GROUP_CONTAINING_VAR_NAME_RE,
) -> str:
    """Generic function for !env yaml tag.

    Deserializes the yaml node tagged with !env, by replacing
    variables names with its values from the process environment.

    Variables should be specified as ${VARIABLE}.
    """
    raise NotImplementedError(
        f"Create implementation of the `env_constructor` "
        f"for node of {type(node)} type."
    )


@env_constructor.register
def _(
    node: yaml.ScalarNode,
    loader: yaml.SafeLoader,
    *,
    string_with_vars_re: re.Pattern = STRING_WITH_VARS_RE,
    group_name: str = GROUP_CONTAINING_VAR_NAME_RE,
) -> str:
    """Implement env_constructor for node: ScalarNode."""

    def get_var(match: re.Match) -> str:
        return os.environ.get(match.group(group_name), "")

    return string_with_vars_re.sub(
        get_var,
        loader.construct_scalar(node),
    )


@env_constructor.register
def _(
    node: yaml.MappingNode,
    loader: yaml.SafeLoader,
    *,
    string_with_vars_re: re.Pattern = STRING_WITH_VARS_RE,
    group_name: str = GROUP_CONTAINING_VAR_NAME_RE,
) -> Dict[str, Any]:
    """Implement env_constructor for node: MappingNode.

    It recursively invokes itself for the nested MappingNode-s.
    """
    return {
        node_key.value: env_constructor(node_value, loader)
        for node_key, node_value in node.value
    }


yaml.add_constructor(
    "!env",
    lambda loader, node: env_constructor(node, loader),
    Loader=yaml.SafeLoader,
)


def read_deploy_config(
    config_path: pathlib.Path,
    _open=open,
) -> Dict[str, str]:
    """Deserialize yaml config into python dict using yaml.SafeLoader.

    Also supports:
    - environment variable expansion

    Use !env tag to indicate the value contains the env var.
    The environment variable should be specified as ${ENV}. For example:

    ```yaml
    cool_new_property: !env ${VAR1} and ${VAR2}
    ```

    Also support marking the nested parts of the yaml, for example:

    ```yaml
    nested_property: !env
      key_0:
        value_0: ${HI_ENV_VAR}
        value_1: ${HI_ENV_VAR}
        value_2: "Just a string"
      key_1: ${HI_ENV_VAR_1}
    ```
    """
    with _open(config_path) as config_io:
        config = yaml.safe_load(config_io)
    return config

test_better_yaml.py

import pytest
from pytest_mock import MockerFixture

from pcty_duck.utils.config_loader import read_deploy_config


@pytest.mark.parametrize(
    (
        "yaml_node_string",
        "env_name",
        "env_value",
        "expected",
    ),
    (
        (
            "!env A string with the ${HI_ENV_VAR}",
            "HI_ENV_VAR",
            "some_value",
            "A string with the some_value",
        ),
        (
            "A string with the ${HI_ENV_VAR}",
            "HI_ENV_VAR",
            "some_value",
            "A string with the ${HI_ENV_VAR}",
        ),
        (
            "!env A string with the ${HI_ENV_VAR}",
            "OTHER_VAR",
            "some_value",
            "A string with the ",
        ),
    ),
)
def test_should_properly_replace_vars_in_flat_structures(
    mocker: MockerFixture,
    monkeypatch,
    yaml_node_string,
    env_name,
    env_value,
    expected,
):
    payload = f"""
    cool_new_property: {yaml_node_string}
    """
    monkeypatch.setenv(env_name, env_value)
    actual = read_deploy_config(
        payload,
        _open=mocker.mock_open(read_data=payload),
    )
    assert actual == {"cool_new_property": expected}


def test_should_work_with_multiple_vars(
    mocker: MockerFixture,
    monkeypatch,
):
    payload = """
    cool_new_property: !env ${VAR1} and ${VAR2}
    """
    monkeypatch.setenv("VAR1", "VAL1")
    monkeypatch.setenv("VAR2", "VAL2")
    actual = read_deploy_config(
        payload,
        _open=mocker.mock_open(read_data=payload),
    )
    assert actual == {"cool_new_property": "VAL1 and VAL2"}


def test_should_work_with_nested_structures(
    mocker: MockerFixture,
    monkeypatch,
):
    monkeypatch.setenv("HI_ENV_VAR", "some_value")
    monkeypatch.setenv("HI_ENV_VAR_1", "some_value_1")
    payload = """
    nested_property: !env
      key_0:
        value_0: ${HI_ENV_VAR}
        value_1: ${HI_ENV_VAR}
        value_2: "Just a string"
      key_1: ${HI_ENV_VAR_1}
    """
    actual = read_deploy_config(
        payload,
        _open=mocker.mock_open(read_data=payload),
    )
    assert actual == {
        "nested_property": {
            "key_0": {
                "value_0": "some_value",
                "value_1": "some_value",
                "value_2": "Just a string",
            },
            "key_1": "some_value_1",
        }
    }