utils.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
工具函数模块
提供通用的工具函数，包括日期处理、文件操作、数据验证等
"""

import os
import logging
import re
import json
from pathlib import Path
from datetime import datetime, timedelta
from typing import Any, List, Dict, Tuple, Optional

# 配置日志
logger = logging.getLogger(__name__)

def ensure_dir_exists(dir_path: str) -> bool:
    """确保目录存在"""
    try:
        Path(dir_path).mkdir(parents=True, exist_ok=True)
        return True
    except Exception as e:
        logger.error(f"创建目录失败：{dir_path} - {e}")
        return False

def load_json_file(file_path: str, default: Any = None) -> Any:
    """加载JSON文件"""
    try:
        if Path(file_path).exists():
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        else:
            logger.warning(f"文件不存在：{file_path}")
            return default
    except Exception as e:
        logger.error(f"加载JSON文件失败：{file_path} - {e}")
        return default

def save_json_file(file_path: str, data: Any) -> bool:
    """保存JSON文件"""
    try:
        ensure_dir_exists(str(Path(file_path).parent))
        
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        return True
    except Exception as e:
        logger.error(f"保存JSON文件失败：{file_path} - {e}")
        return False

def load_text_file(file_path: str) -> str:
    """加载文本文件"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        logger.error(f"加载文本文件失败：{file_path} - {e}")
        return ""

def save_text_file(file_path: str, content: str) -> bool:
    """保存文本文件"""
    try:
        ensure_dir_exists(str(Path(file_path).parent))
        
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)
        
        return True
    except Exception as e:
        logger.error(f"保存文本文件失败：{file_path} - {e}")
        return False

def format_date(date_str: str, from_format: str = "%Y%m%d", to_format: str = "%Y-%m-%d") -> str:
    """格式化日期字符串"""
    try:
        date_obj = datetime.strptime(date_str, from_format)
        return date_obj.strftime(to_format)
    except Exception as e:
        logger.error(f"日期格式化失败：{date_str} - {e}")
        return date_str

def get_trade_date(date_obj: Optional[datetime] = None) -> str:
    """获取交易日日期字符串"""
    try:
        if date_obj is None:
            date_obj = datetime.now()
        
        return date_obj.strftime("%Y%m%d")
    except Exception as e:
        logger.error(f"获取交易日日期失败：{e}")
        return datetime.now().strftime("%Y%m%d")

def get_prev_trading_date(trade_date: str, num_days: int = 1) -> str:
    """获取前N个交易日的日期"""
    try:
        date_obj = datetime.strptime(trade_date, "%Y%m%d")
        prev_date = date_obj - timedelta(days=num_days)
        return prev_date.strftime("%Y%m%d")
    except Exception as e:
        logger.error(f"获取前交易日日期失败：{trade_date} - {e}")
        return trade_date

def validate_stock_code(stock_code: str) -> bool:
    """验证股票代码格式"""
    # 股票代码格式：6位数字 + 市场后缀（.SH/.SZ）
    pattern = r"^\d{6}\.(SH|SZ)$"
    return bool(re.match(pattern, stock_code))

def validate_date_string(date_str: str, format: str = "%Y%m%d") -> bool:
    """验证日期字符串格式"""
    try:
        datetime.strptime(date_str, format)
        return True
    except:
        return False

def calculate_volatility(prices: List[float]) -> float:
    """计算价格波动率（标准差）"""
    try:
        import numpy as np
        return np.std(prices)
    except Exception as e:
        logger.error(f"计算波动率失败：{e}")
        return 0.0

def calculate_moving_average(prices: List[float], period: int) -> List[float]:
    """计算移动平均"""
    try:
        import numpy as np
        return np.convolve(prices, np.ones(period)/period, mode='valid').tolist()
    except Exception as e:
        logger.error(f"计算移动平均失败：{e}")
        return []

def calculate_rsi(prices: List[float], period: int = 14) -> List[float]:
    """计算RSI指标"""
    try:
        import numpy as np
        
        deltas = np.diff(prices)
        gains = np.where(deltas > 0, deltas, 0)
        losses = np.where(deltas < 0, -deltas, 0)
        
        avg_gain = np.convolve(gains, np.ones(period)/period, mode='valid')
        avg_loss = np.convolve(losses, np.ones(period)/period, mode='valid')
        
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
        
        return rsi.tolist()
    except Exception as e:
        logger.error(f"计算RSI失败：{e}")
        return []

def calculate_macd(prices: List[float], fast: int = 12, slow: int = 26, signal: int = 9) -> Tuple[List[float], List[float], List[float]]:
    """计算MACD指标"""
    try:
        import numpy as np
        
        ema_fast = calculate_exponential_moving_average(prices, fast)
        ema_slow = calculate_exponential_moving_average(prices, slow)
        
        macd = np.array(ema_fast) - np.array(ema_slow)
        signal_line = calculate_exponential_moving_average(macd, signal)
        histogram = macd - signal_line
        
        return macd.tolist(), signal_line.tolist(), histogram.tolist()
    except Exception as e:
        logger.error(f"计算MACD失败：{e}")
        return [], [], []

def calculate_exponential_moving_average(prices: List[float], period: int) -> List[float]:
    """计算指数移动平均"""
    try:
        import numpy as np
        
        alpha = 2 / (period + 1)
        ema = []
        
        for i, price in enumerate(prices):
            if i == 0:
                ema.append(price)
            else:
                ema.append(alpha * price + (1 - alpha) * ema[i-1])
        
        return ema
    except Exception as e:
        logger.error(f"计算指数移动平均失败：{e}")
        return []

def format_number(number: float, format_str: str = "{:.2f}") -> str:
    """格式化数字"""
    try:
        return format_str.format(number)
    except Exception as e:
        logger.error(f"数字格式化失败：{number} - {e}")
        return str(number)

def truncate_text(text: str, max_length: int, suffix: str = "...") -> str:
    """截断文本"""
    if len(text) <= max_length:
        return text
    return text[:max_length - len(suffix)] + suffix

def remove_duplicates(lst: List[Any]) -> List[Any]:
    """去除列表中的重复元素"""
    seen = []
    result = []
    
    for item in lst:
        if item not in seen:
            seen.append(item)
            result.append(item)
    
    return result

def get_file_size(file_path: str) -> int:
    """获取文件大小（字节）"""
    try:
        return os.path.getsize(file_path)
    except Exception as e:
        logger.error(f"获取文件大小失败：{file_path} - {e}")
        return 0

def get_file_modification_time(file_path: str) -> datetime:
    """获取文件修改时间"""
    try:
        return datetime.fromtimestamp(os.path.getmtime(file_path))
    except Exception as e:
        logger.error(f"获取文件修改时间失败：{file_path} - {e}")
        return datetime.now()

def split_list(lst: List[Any], n: int) -> List[List[Any]]:
    """将列表分成N个大小相等的子列表"""
    try:
        chunk_size = len(lst) // n
        remainder = len(lst) % n
        
        result = []
        start = 0
        
        for i in range(n):
            end = start + chunk_size + (1 if i < remainder else 0)
            result.append(lst[start:end])
            start = end
        
        return result
    except Exception as e:
        logger.error(f"列表分割失败：{e}")
        return [lst]

def merge_dicts(*dicts: Dict[Any, Any]) -> Dict[Any, Any]:
    """合并多个字典"""
    result = {}
    
    for d in dicts:
        result.update(d)
    
    return result

def safe_divide(a: float, b: float, default: float = 0.0) -> float:
    """安全除法"""
    try:
        if b == 0:
            return default
        return a / b
    except Exception as e:
        logger.error(f"除法运算失败：{a} / {b} - {e}")
        return default

def retry_on_exception(max_retries: int = 3, delay: int = 1):
    """异常重试装饰器"""
    import time
    from functools import wraps
    
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            
            for i in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    logger.warning(f"执行失败，第 {i+1} 次重试：{e}")
                    time.sleep(delay)
            
            logger.error(f"执行失败，已达到最大重试次数：{last_exception}")
            raise last_exception
        
        return wrapper
    return decorator