Last active
April 20, 2023 19:50
-
-
Save zhukovgreen/c4b4d785d48b5be572ec8a4573f1b778 to your computer and use it in GitHub Desktop.
yaml with env vars
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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", | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment