关于python:python-解析-yaml-环境变量替换-value

import os
import re
from functools import partial

import yaml


has_regex_module = False
ENV_VAR_MATCHER = re.compile(
    r"""
        \$\{       # match characters `${` literally
        ([^}:\s]+) # 1st group: matches any character except `}` or `:`
        :?         # matches the literal `:` character zero or one times
        ([^}]+)?   # 2nd group: matches any character except `}`
        \}         # match character `}` literally
    """, re.VERBOSE
)


IMPLICIT_ENV_VAR_MATCHER = re.compile(
    r"""
        .*          # matches any number of any characters
        \$\{.*\}    # matches any number of any characters
                    # between `${` and `}` literally
        .*          # matches any number of any characters
    """, re.VERBOSE
)


RECURSIVE_ENV_VAR_MATCHER = re.compile(
    r"""
        \$\{       # match characters `${` literally
        ([^}]+)?   # matches any character except `}`
        \}         # match character `}` literally
        ([^$}]+)?  # matches any character except `}` or `$`
        \}         # match character `}` literally
    """,
    re.VERBOSE,
)


def _replace_env_var(match):
    env_var, default = match.groups()
    value = os.environ.get(env_var, None)
    if value is None:
        # expand default using other vars
        if default is None:
            # regex module return None instead of
            #  ''if engine didn't entered default capture group
            default = ''

        value = default
        while IMPLICIT_ENV_VAR_MATCHER.match(value):  # pragma: no cover
            value = ENV_VAR_MATCHER.sub(_replace_env_var, value)
    return value


def env_var_constructor(loader, node, raw=False):
    raw_value = loader.construct_scalar(node)

    # detect and error on recursive environment variables
    if not has_regex_module and RECURSIVE_ENV_VAR_MATCHER.match(raw_value):  # pragma: no cover
        raise Exception("Nested environment variable lookup requires the `regex` module")
    value = ENV_VAR_MATCHER.sub(_replace_env_var, raw_value)
    if value == raw_value:
        return value  # avoid recursion
    return value if raw else yaml.safe_load(value)


def setup_yaml_parser():
    yaml.add_constructor('!env_var', env_var_constructor, yaml.SafeLoader)
    yaml.add_constructor(
        '!raw_env_var',
        partial(env_var_constructor, raw=True),
        yaml.SafeLoader
    )
    yaml.add_implicit_resolver('!env_var', IMPLICIT_ENV_VAR_MATCHER, Loader=yaml.SafeLoader)

# copy from nameko https://github.com/nameko/nameko/blob/v2.14.1/nameko/cli/main.py

写个单元测试

# coding=utf-8
import csv
import json
import os
import time
import unittest
from collections import namedtuple
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from uuid import uuid4
# from elasticsearch import Elasticsearch
from loguru import logger
import settings
from mark import BASE_DIR
from pathlib import Path
import yaml

TESTING_BASE_DIR = Path(__file__).resolve().parent


class TestEnvParse(unittest.TestCase):
    def test_parse_yaml_with_env(self):
        """
        测试 es 是否联通
        python -m unittest testing.test_env.TestEnvParse.test_parse_yaml_with_env
        """
        from pon.events import EventletEventRunner

        config_filepath = TESTING_BASE_DIR/'config.yaml'

        with open(config_filepath, 'r', encoding='utf-8') as f:
            config: dict[str, dict] = yaml.safe_load(f)

        logger.debug(config)

筹备的 yaml 文件

AMQP_URI: amqp://${RABBIT_USER:guest}:${RABBIT_PASSWORD:guest}@${RABBIT_HOST:localhost}:${RABBIT_PORT:5672}/${RABBIT_VHOST:/}

prd:
    database:
        mysql:
            host: ${MYSQL_HOST:110.110.100.110}
            port: 3306
            username: you
            password: you
            database_name: haha

        elasticsearch:
            host: 1110.110.110.100
            port: 9200
            username: hah
            password: hah
            index_name: hehe

运行命令:MYSQL_HOST=192.168.31.245 python -m unittest testing.test_env.TestEnvParse.test_parse_yaml_with_env

2022-09-27 13:54:29.565 | DEBUG    | testing.test_env:test_parse_yaml_with_env:33 - {'AMQP_URI': 'amqp://pon:[email protected]:5672//', 'prd': {'database': {'mysql': {'host': '192.168.31.245', 'port': 3306, 'username': 'you', 'password': 'you', 'database_name': 'haha'}, 'elasticsearch': {'host': '1110.110.110.100', 'port': 9200, 'username': 'hah', 'password': 'hah', 'index_name': 'hehe'}}}}

运行后果，能够看到 mysql 的 host 曾经被环境变量替换了