spider_work/pcg_spider/YamlLoader.py

# -*- coding: utf-8 -*-
# Author : Charley
# Python : 3.10.8
# Date   : 2025/12/22 10:44
import os, re
import yaml

regex = re.compile(r'^\$\{(?P<ENV>[A-Z_\-]+:)?(?P<VAL>[\w.]+)}$')


class YamlConfig:
    def __init__(self, config):
        self.config = config

    def get(self, key: str):
        return YamlConfig(self.config.get(key))

    def getValueAsString(self, key: str):
        try:
            match = regex.match(self.config[key])
            group = match.groupdict()
            if group['ENV'] is not None:
                env = group['ENV'][:-1]
                return os.getenv(env, group['VAL'])
            return None
        except:
            return self.config[key]

    def getValueAsInt(self, key: str):
        try:
            match = regex.match(self.config[key])
            group = match.groupdict()
            if group['ENV'] is not None:
                env = group['ENV'][:-1]
                return int(os.getenv(env, group['VAL']))
            return 0
        except:
            return int(self.config[key])

    def getValueAsBool(self, key: str):
        try:
            match = regex.match(self.config[key])
            group = match.groupdict()
            if group['ENV'] is not None:
                env = group['ENV'][:-1]
                return bool(os.getenv(env, group['VAL']))
            return False
        except:
            return bool(self.config[key])


def _resolve_path(path: str) -> str:
    """
    解析 yaml 文件路径，按优先级查找：
      1) 绝对路径或 cwd 下存在 → 直接用（保留旧行为，向后兼容）
      2) 调用方主脚本所在目录 → 兜底，方便打包后从任意 cwd 启动
    :param path: (str) 用户传入的路径，默认 'application.yml'
    :return: (str) 实际可读取的完整路径；找不到则返回原 path 让 open() 抛错
    """
    # 1) 旧行为：cwd 或绝对路径
    if os.path.exists(path):
        return path

    # 2) 主脚本目录（__main__.__file__）
    try:
        import __main__
        main_file = getattr(__main__, '__file__', None)
        if main_file:
            candidate = os.path.join(os.path.dirname(os.path.abspath(main_file)), path)
            if os.path.exists(candidate):
                return candidate
    except Exception:
        pass

    return path


def readYaml(path: str = 'application.yml', profile: str = None) -> YamlConfig:
    """
    读取 yaml 配置。
    :param path: (str) yaml 文件路径，默认 'application.yml'。
                       优先 cwd / 绝对路径（保留旧行为），找不到再 fallback 到主脚本所在目录。
    :param profile: (str) 可选环境后缀，如 'dev' 会额外加载 'application-dev.yml' 并 update
    :return: (YamlConfig) 配置访问对象
    :raises FileNotFoundError: cwd 和主脚本目录都找不到时抛出
    """
    real_path = _resolve_path(path)
    with open(real_path, encoding='utf-8') as fd:
        conf = yaml.load(fd, Loader=yaml.FullLoader)

    if profile is not None:
        result = real_path.rsplit('.', 1)
        profiledYaml = f'{result[0]}-{profile}.{result[1]}'
        if os.path.exists(profiledYaml):
            with open(profiledYaml, encoding='utf-8') as fd:
                conf.update(yaml.load(fd, Loader=yaml.FullLoader))

    return YamlConfig(conf)