Skip to content

Commit

Permalink
custom env loader
Browse files Browse the repository at this point in the history
  • Loading branch information
MDobransky committed Sep 5, 2024
1 parent 6b2831d commit 6ca429d
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 2 deletions.
1 change: 1 addition & 0 deletions .flake8
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,4 @@ extend-ignore =
D100,
D104,
D107,
E203,
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ pandas = "^2.1.0"
flake8-broken-line = "^1.0.0"
loguru = "^0.7.2"
importlib-metadata = "^7.2.1"
env_yaml = "^0.0.3"

[tool.poetry.dev-dependencies]
pyspark = "^3.4.1"
Expand Down
28 changes: 28 additions & 0 deletions rialto/common/env_yaml.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import os
import re

import yaml
from loguru import logger

__all__ = ["EnvLoader"]

_path_matcher = re.compile(r"\$\{(?P<env_name>[^}^{:]+)(?::(?P<default_value>[^}^{]*))?\}")


def _path_constructor(loader, node):
value = node.value
match = _path_matcher.match(value)
sub = os.getenv(match.group("env_name"), match.group("default_value"))
new_value = value[0 : match.start()] + sub + value[match.end() :]
logger.info(f"Config: Replacing {value}, with {new_value}")
return new_value


class EnvLoader(yaml.SafeLoader):
"""Custom loader that replaces values with environment variables"""

pass


EnvLoader.add_implicit_resolver("!env_substitute", _path_matcher, None)
EnvLoader.add_constructor("!env_substitute", _path_constructor)
3 changes: 2 additions & 1 deletion rialto/common/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,11 @@

import pyspark.sql.functions as F
import yaml
from env_yaml import EnvLoader
from pyspark.sql import DataFrame
from pyspark.sql.types import FloatType

from rialto.common.env_yaml import EnvLoader


def load_yaml(path: str) -> Any:
"""
Expand Down

0 comments on commit 6ca429d

Please sign in to comment.