230 lines
7.4 KiB
Python
230 lines
7.4 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
try:
|
||
import yaml
|
||
YAML_AVAILABLE = True
|
||
except ImportError:
|
||
YAML_AVAILABLE = False
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
CONFIG_PATH = ROOT / "config.yaml"
|
||
LOCAL_CONFIG_PATH = ROOT / "config.local.yaml"
|
||
|
||
# 嵌套 YAML 路径 → 旧版扁平静态 key 映射表
|
||
# 确保从 config.yaml 迁移后，旧代码仍能按 flat key 取值
|
||
_NESTED_TO_FLAT_MAP = {
|
||
# output
|
||
("output", "base_dir"): "OUTPUT_BASE_DIR",
|
||
# device
|
||
("device", "is_emulator"): "IS_EMULATOR",
|
||
("device", "physical_ip"): "PHYSICAL_DEVICE_IP",
|
||
("device", "physical_adb_port"): "PHYSICAL_ADB_PORT",
|
||
("device", "physical_serial"): "PHYSICAL_DEVICE_SERIAL",
|
||
("device", "airtest_adb_host"): "AIRTEST_ADB_HOST",
|
||
# app
|
||
("app", "keep_app"): "KEEP_APP",
|
||
# policy
|
||
("policy", "name"): "POLICY",
|
||
("policy", "enable_guiagent"): "ENABLE_GUIAGENT",
|
||
("policy", "cv_mode"): "CV_MODE",
|
||
# execution
|
||
("execution", "timeout"): "TIMEOUT",
|
||
("execution", "block_timeout"): "BLOCK_TIMEOUT",
|
||
# model
|
||
("model", "current_name"): "CURRENT_MODEL_NAME",
|
||
# logging
|
||
("logging", "level"): "LOG_LEVEL",
|
||
# redis
|
||
("redis", "host"): "CONTROL_REDIS_HOST",
|
||
("redis", "port"): "CONTROL_REDIS_PORT",
|
||
("redis", "db"): "CONTROL_REDIS_DB",
|
||
("redis", "channel_namespace"): "CONTROL_CHANNEL_NAMESPACE",
|
||
# shared_paths
|
||
("shared_paths", "traffic_data"): "TRAFFIC_DATA_SHARE",
|
||
("shared_paths", "traversal_log"): "TRAVERSAL_LOG_SHARE",
|
||
("shared_paths", "clean_backup"): "CLEAN_BACKUP_PATH",
|
||
("shared_paths", "local_apk"): "LOCAL_APK_PATH",
|
||
("shared_paths", "apkpure_apk"): "APKPURE_APK_PATH",
|
||
("shared_paths", "app_block_app_list"): "APP_BLOCK_APP_LIST",
|
||
("shared_paths", "app_block_url_lib"): "APP_BLOCK_URL_LIB",
|
||
# git
|
||
("git", "repo_url"): "GIT_REPO_URL",
|
||
# notifications
|
||
("notifications", "wechat"): "WECHAT_TOKENS",
|
||
("notifications", "wecom"): "WECOM_TOKENS",
|
||
# mumu
|
||
("mumu", "manager_path"): "MUMU_MANAGER_PATH",
|
||
("mumu", "local_import_dir"): "LOCAL_IMPORT_DIR",
|
||
("mumu", "default_vm_index"): "DEFAULT_VM_INDEX",
|
||
("mumu", "max_vm_index"): "MAX_VM_INDEX",
|
||
# network
|
||
("network", "bridge_card"): "NET_BRIDGE_CARD",
|
||
# llm_api
|
||
("llm_api", "key_pool_url"): "KEY_POOL_URL",
|
||
}
|
||
|
||
|
||
def _normalize_config(config: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""将嵌套 YAML 配置归一化为兼容旧版 flat key 格式。
|
||
|
||
如果配置已包含旧版 flat key（如 CONTROL_REDIS_HOST），
|
||
说明是旧格式 JSON，不做转换直接返回。
|
||
否则按照 _NESTED_TO_FLAT_MAP 将嵌套路径值注入到顶层。
|
||
"""
|
||
# 已包含 flat key 说明是旧格式，跳过
|
||
if "CONTROL_REDIS_HOST" in config:
|
||
return config
|
||
|
||
normalized = dict(config)
|
||
for nested_keys, flat_key in _NESTED_TO_FLAT_MAP.items():
|
||
if flat_key in normalized:
|
||
continue
|
||
value = config
|
||
for key in nested_keys:
|
||
if isinstance(value, dict):
|
||
value = value.get(key)
|
||
else:
|
||
value = None
|
||
break
|
||
if value is not None:
|
||
normalized[flat_key] = value
|
||
|
||
return normalized
|
||
|
||
|
||
def _expand_env_vars(value: Any) -> Any:
|
||
"""递归展开配置中的环境变量引用，格式：${ENV_VAR_NAME}"""
|
||
if isinstance(value, str):
|
||
# 匹配 ${VAR_NAME} 格式
|
||
pattern = re.compile(r'\$\{([^}]+)\}')
|
||
|
||
def replacer(match):
|
||
var_name = match.group(1)
|
||
env_value = os.environ.get(var_name)
|
||
if env_value is None:
|
||
# 环境变量不存在时保持原样，便于调试
|
||
return match.group(0)
|
||
return env_value
|
||
|
||
return pattern.sub(replacer, value)
|
||
elif isinstance(value, dict):
|
||
return {k: _expand_env_vars(v) for k, v in value.items()}
|
||
elif isinstance(value, list):
|
||
return [_expand_env_vars(item) for item in value]
|
||
return value
|
||
|
||
|
||
def _load_config_dict(path: Path, *, required: bool) -> Dict[str, Any]:
|
||
"""加载配置文件（支持 JSON 和 YAML 格式）
|
||
|
||
Args:
|
||
path: 配置文件路径
|
||
required: 是否必须存在，不存在时是否抛出异常
|
||
|
||
Returns:
|
||
配置字典
|
||
|
||
Raises:
|
||
FileNotFoundError: 当 required=True 且文件不存在时
|
||
ImportError: 加载 YAML 文件但 PyYAML 未安装时
|
||
ValueError: 配置文件格式不正确时
|
||
"""
|
||
if not path.exists():
|
||
if required:
|
||
raise FileNotFoundError(
|
||
f"Configuration file not found: {path}\n"
|
||
f"Hint: Make sure the config file exists at the expected location."
|
||
)
|
||
return {}
|
||
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
if path.suffix in {".yaml", ".yml"}:
|
||
if not YAML_AVAILABLE:
|
||
raise ImportError(
|
||
f"PyYAML is required to load {path.name}.\n"
|
||
f"Install it with: pip install pyyaml"
|
||
)
|
||
payload = yaml.safe_load(handle)
|
||
elif path.suffix == ".json":
|
||
payload = json.load(handle)
|
||
else:
|
||
raise ValueError(f"Unsupported config format: {path.suffix} (supported: .json, .yaml, .yml)")
|
||
|
||
if not isinstance(payload, dict):
|
||
raise ValueError(f"Config payload must be a JSON/YAML object (dict), got {type(payload).__name__}: {path}")
|
||
|
||
# 展开环境变量引用
|
||
return _expand_env_vars(payload)
|
||
|
||
|
||
def _merge_dict(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
|
||
merged = dict(base)
|
||
for key, value in override.items():
|
||
current = merged.get(key)
|
||
if isinstance(current, dict) and isinstance(value, dict):
|
||
merged[key] = _merge_dict(current, value)
|
||
continue
|
||
merged[key] = value
|
||
return merged
|
||
|
||
|
||
def _load_local_config() -> Dict[str, Any]:
|
||
"""加载本地配置文件（config.local.yaml）
|
||
|
||
Returns:
|
||
本地配置字典，如果文件不存在则返回空字典
|
||
"""
|
||
if LOCAL_CONFIG_PATH.exists():
|
||
logger.debug(f"Loading local config from: {LOCAL_CONFIG_PATH}")
|
||
return _load_config_dict(LOCAL_CONFIG_PATH, required=False)
|
||
|
||
logger.debug("No local config file found (optional)")
|
||
return {}
|
||
|
||
|
||
def load_layered_config() -> Dict[str, Any]:
|
||
"""加载分层配置（主配置 + 本地覆盖配置）
|
||
|
||
配置加载顺序：
|
||
1. config.yaml 作为基础配置
|
||
2. config.local.yaml 覆盖基础配置（如果存在）
|
||
3. 归一化为兼容旧版 flat key 格式
|
||
|
||
Returns:
|
||
合并并归一化后的配置字典
|
||
|
||
Raises:
|
||
FileNotFoundError: config.yaml 不存在时
|
||
"""
|
||
config = _load_config_dict(CONFIG_PATH, required=True)
|
||
local_config = _load_local_config()
|
||
config = _merge_dict(config, local_config)
|
||
return _normalize_config(config)
|
||
|
||
|
||
def load_config(yaml_path: Optional[str] = None) -> Dict[str, Any]:
|
||
"""加载配置文件
|
||
|
||
Args:
|
||
yaml_path: 自定义配置文件路径，None 时使用默认分层配置
|
||
|
||
Returns:
|
||
归一化后的配置字典（同时包含嵌套结构和旧版 flat key）
|
||
|
||
Raises:
|
||
FileNotFoundError: 配置文件不存在时
|
||
"""
|
||
if yaml_path:
|
||
return _normalize_config(_load_config_dict(Path(yaml_path).resolve(), required=True))
|
||
return load_layered_config()
|