File size: 5,580 Bytes
5669b22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | # config_manager/utils.py
import yaml
from pathlib import Path
from typing import Union, Dict, Any, TypeVar
from pydantic import BaseModel, ValidationError
import os
import re
import chardet
from loguru import logger
from .main import Config
T = TypeVar("T", bound=BaseModel)
def read_yaml(config_path: str) -> Dict[str, Any]:
"""
Read the specified YAML configuration file with environment variable substitution
and guess encoding. Return the configuration data as a dictionary.
Args:
config_path: Path to the YAML configuration file.
Returns:
Configuration data as a dictionary.
Raises:
FileNotFoundError: If the configuration file is not found.
IOError: If the configuration file cannot be read.
"""
if not os.path.exists(config_path):
raise FileNotFoundError(f"Configuration file not found: {config_path}")
content = load_text_file_with_guess_encoding(config_path)
if not content:
raise IOError(f"Failed to read configuration file: {config_path}")
# Replace environment variables
pattern = re.compile(r"\$\{(\w+)\}")
def replacer(match):
env_var = match.group(1)
return os.getenv(env_var, match.group(0))
content = pattern.sub(replacer, content)
try:
return yaml.safe_load(content)
except yaml.YAMLError as e:
logger.critical(f"Error parsing YAML file: {e}")
raise e
def validate_config(config_data: dict) -> Config:
"""
Validate configuration data against the Config model.
Args:
config_data: Configuration data to validate.
Returns:
Validated Config object.
Raises:
ValidationError: If the configuration fails validation.
"""
try:
return Config(**config_data)
except ValidationError as e:
logger.critical(f"Error validating configuration: {e}")
logger.error("Configuration data:")
logger.error(config_data)
raise e
def load_text_file_with_guess_encoding(file_path: str) -> str | None:
"""
Load a text file with guessed encoding.
Parameters:
- file_path (str): The path to the text file.
Returns:
- str: The content of the text file or None if an error occurred.
"""
encodings = ["utf-8", "utf-8-sig", "gbk", "gb2312", "ascii", "cp936"]
for encoding in encodings:
try:
with open(file_path, "r", encoding=encoding) as file:
return file.read()
except UnicodeDecodeError:
continue
# If common encodings fail, try chardet to guess the encoding
try:
with open(file_path, "rb") as file:
raw_data = file.read()
detected = chardet.detect(raw_data)
if detected["encoding"]:
return raw_data.decode(detected["encoding"])
except Exception as e:
logger.error(f"Error detecting encoding for config file {file_path}: {e}")
return None
def save_config(config: BaseModel, config_path: Union[str, Path]):
"""
Saves a Pydantic model to a YAML configuration file.
Args:
config: The Pydantic model to save.
config_path: Path to the YAML configuration file.
"""
config_file = Path(config_path)
config_data = config.model_dump(
by_alias=True, exclude_unset=True, exclude_none=True
)
try:
with open(config_file, "w", encoding="utf-8") as f:
yaml.dump(config_data, f, allow_unicode=True)
except yaml.YAMLError as e:
raise yaml.YAMLError(f"Error writing YAML file: {e}")
def scan_config_alts_directory(config_alts_dir: str) -> list[dict]:
"""
Scan the config_alts directory and return a list of config information.
Each config info contains the filename and its display name from the config.
Parameters:
- config_alts_dir (str): The path to the config_alts directory.
Returns:
- list[dict]: A list of dicts containing config info:
- filename: The actual config file name
- name: Display name from config, falls back to filename if not specified
"""
config_files = []
# Add default config first
default_config = read_yaml("conf.yaml")
config_files.append(
{
"filename": "conf.yaml",
"name": default_config.get("character_config", {}).get(
"conf_name", "conf.yaml"
)
if default_config
else "conf.yaml",
}
)
# Scan other configs
for root, _, files in os.walk(config_alts_dir):
for file in files:
if file.endswith(".yaml"):
config: dict = read_yaml(os.path.join(root, file))
config_files.append(
{
"filename": file,
"name": config.get("character_config", {}).get(
"conf_name", file
)
if config
else file,
}
)
logger.debug(f"Found config files: {config_files}")
return config_files
def scan_bg_directory() -> list[str]:
bg_files = []
bg_dir = "backgrounds"
for root, _, files in os.walk(bg_dir):
for file in files:
if file.endswith((".jpg", ".jpeg", ".png", ".gif")):
bg_files.append(file)
return bg_files
|