import osimport refrom functools import partialimport yamlhas_regex_module = FalseENV_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 valuedef 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-8import csvimport jsonimport osimport timeimport unittestfrom collections import namedtuplefrom concurrent.futures import ThreadPoolExecutorfrom pathlib import Pathfrom uuid import uuid4# from elasticsearch import Elasticsearchfrom loguru import loggerimport settingsfrom mark import BASE_DIRfrom pathlib import Pathimport yamlTESTING_BASE_DIR = Path(__file__).resolve().parentclass 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 曾经被环境变量替换了