DualFlow/common/utils/config.py

"""
配置管理工具
"""
import yaml
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional, Union
import logging

logger = logging.getLogger(__name__)


class Config:
    """配置管理类"""
    
    def __init__(self, config_path: Optional[Union[str, Path]] = None):
        """初始化配置"""
        self.config = {}
        if config_path:
            self.load_config(config_path)
    
    def load_config(self, config_path: Union[str, Path]) -> None:
        """加载配置文件"""
        config_path = Path(config_path)
        
        if not config_path.exists():
            raise FileNotFoundError(f"配置文件不存在: {config_path}")
        
        if config_path.suffix.lower() == '.yaml' or config_path.suffix.lower() == '.yml':
            self.load_yaml(config_path)
        elif config_path.suffix.lower() == '.json':
            self.load_json(config_path)
        else:
            raise ValueError(f"不支持的配置文件格式: {config_path.suffix}")
    
    def load_yaml(self, config_path: Path) -> None:
        """加载YAML配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                self.config = yaml.safe_load(f)
            logger.info(f"成功加载YAML配置文件: {config_path}")
        except Exception as e:
            logger.error(f"加载YAML配置文件失败: {e}")
            raise
    
    def load_json(self, config_path: Path) -> None:
        """加载JSON配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                self.config = json.load(f)
            logger.info(f"成功加载JSON配置文件: {config_path}")
        except Exception as e:
            logger.error(f"加载JSON配置文件失败: {e}")
            raise
    
    def get(self, key: str, default: Any = None) -> Any:
        """获取配置值"""
        keys = key.split('.')
        value = self.config
        
        for k in keys:
            if isinstance(value, dict) and k in value:
                value = value[k]
            else:
                return default
        
        return value
    
    def set(self, key: str, value: Any) -> None:
        """设置配置值"""
        keys = key.split('.')
        config = self.config
        
        for k in keys[:-1]:
            if k not in config:
                config[k] = {}
            config = config[k]
        
        config[keys[-1]] = value
    
    def update(self, other_config: Dict[str, Any]) -> None:
        """更新配置"""
        self.config.update(other_config)
    
    def save_yaml(self, config_path: Union[str, Path]) -> None:
        """保存为YAML文件"""
        config_path = Path(config_path)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            with open(config_path, 'w', encoding='utf-8') as f:
                yaml.dump(self.config, f, default_flow_style=False, allow_unicode=True)
            logger.info(f"成功保存YAML配置文件: {config_path}")
        except Exception as e:
            logger.error(f"保存YAML配置文件失败: {e}")
            raise
    
    def save_json(self, config_path: Union[str, Path]) -> None:
        """保存为JSON文件"""
        config_path = Path(config_path)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, indent=2, ensure_ascii=False)
            logger.info(f"成功保存JSON配置文件: {config_path}")
        except Exception as e:
            logger.error(f"保存JSON配置文件失败: {e}")
            raise
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return self.config.copy()
    
    def __getitem__(self, key: str) -> Any:
        """支持字典式访问"""
        return self.get(key)
    
    def __setitem__(self, key: str, value: Any) -> None:
        """支持字典式设置"""
        self.set(key, value)
    
    def __contains__(self, key: str) -> bool:
        """支持in操作符"""
        return self.get(key) is not None


class EnvironmentConfig:
    """环境变量配置"""
    
    def __init__(self, prefix: str = ""):
        """初始化环境配置"""
        self.prefix = prefix.upper()
    
    def get(self, key: str, default: Any = None, type_func: type = str) -> Any:
        """获取环境变量"""
        env_key = f"{self.prefix}_{key.upper()}" if self.prefix else key.upper()
        value = os.getenv(env_key, default)
        
        if value is None:
            return default
        
        try:
            return type_func(value)
        except (ValueError, TypeError):
            logger.warning(f"无法转换环境变量 {env_key} 为 {type_func.__name__}, 使用默认值: {default}")
            return default
    
    def get_bool(self, key: str, default: bool = False) -> bool:
        """获取布尔环境变量"""
        value = self.get(key, str(default))
        return value.lower() in ('true', '1', 'yes', 'on')
    
    def get_int(self, key: str, default: int = 0) -> int:
        """获取整数环境变量"""
        return self.get(key, default, int)
    
    def get_float(self, key: str, default: float = 0.0) -> float:
        """获取浮点数环境变量"""
        return self.get(key, default, float)
    
    def get_list(self, key: str, default: list = None, separator: str = ',') -> list:
        """获取列表环境变量"""
        if default is None:
            default = []
        
        value = self.get(key, "")
        if not value:
            return default
        
        return [item.strip() for item in value.split(separator)]


def load_config_from_env(config_class: type, env_prefix: str = "") -> Any:
    """从环境变量加载配置类"""
    env_config = EnvironmentConfig(env_prefix)
    
    # 获取配置类的所有属性
    config_instance = config_class()
    
    for attr_name in dir(config_instance):
        if not attr_name.startswith('_'):
            attr_value = getattr(config_instance, attr_name)
            if not callable(attr_value):
                # 从环境变量获取值
                env_value = env_config.get(attr_name, attr_value)
                setattr(config_instance, attr_name, env_value)
    
    return config_instance