utils.py

"""
BaoStock数据缓存系统 - 工具函数模块
包含各种工具函数
"""

from datetime import datetime, timedelta
import pandas as pd
import numpy as np


def format_date(date_str):
    """
    格式化日期
    """
    if not date_str:
        return ""
    
    try:
        # 尝试不同的日期格式
        for fmt in ['%Y-%m-%d', '%Y%m%d', '%Y/%m/%d']:
            try:
                dt = datetime.strptime(date_str, fmt)
                return dt.strftime('%Y-%m-%d')
            except ValueError:
                continue
    except Exception as e:
        pass
    
    return date_str


def get_date_range(days=30):
    """
    获取日期范围
    """
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days)
    return start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d')


def validate_stock_code(code):
    """
    验证股票代码
    """
    if not code:
        return False
    
    # 检查是否为正确的格式
    if '.' in code:
        parts = code.split('.')
        if len(parts) == 2:
            exchange, symbol = parts
            if exchange in ['sh', 'sz'] and len(symbol) == 6 and symbol.isdigit():
                return True
    
    return False


def convert_to_numeric(df, columns=None):
    """
    将DataFrame列转换为数值类型
    """
    if df is None or df.empty:
        return df
    
    if columns is None:
        columns = ['open', 'high', 'low', 'close', 'volume', 'amount']
    
    for col in columns:
        if col in df.columns:
            try:
                df[col] = pd.to_numeric(df[col])
            except Exception as e:
                pass
    
    return df


def calculate_returns(df, price_col='close'):
    """
    计算收益率
    """
    if df is None or df.empty or price_col not in df.columns:
        return df
    
    try:
        df['return'] = df[price_col].astype(float).pct_change()
    except Exception as e:
        pass
    
    return df


def calculate_indicators(df):
    """
    计算基本技术指标
    """
    if df is None or df.empty:
        return df
    
    df = convert_to_numeric(df)
    
    # 计算移动平均线
    if 'close' in df.columns:
        df['ma5'] = df['close'].rolling(window=5).mean()
        df['ma10'] = df['close'].rolling(window=10).mean()
        df['ma20'] = df['close'].rolling(window=20).mean()
    
    # 计算成交量移动平均
    if 'volume' in df.columns:
        df['ma_volume5'] = df['volume'].rolling(window=5).mean()
    
    return df


def filter_trading_days(df, date_col='date'):
    """
    过滤交易日
    """
    if df is None or df.empty or date_col not in df.columns:
        return df
    
    # 移除周末
    df[date_col] = pd.to_datetime(df[date_col])
    df = df[df[date_col].dt.dayofweek < 5]
    
    return df


def resample_data(df, freq='D', date_col='date', price_col='close'):
    """
    重采样数据
    """
    if df is None or df.empty:
        return df
    
    df[date_col] = pd.to_datetime(df[date_col])
    df.set_index(date_col, inplace=True)
    
    # 重采样
    resampled = df.resample(freq).agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last',
        'volume': 'sum',
        'amount': 'sum'
    }).dropna()
    
    resampled.reset_index(inplace=True)
    return resampled


def split_data(df, train_ratio=0.8):
    """
    分割训练集和测试集
    """
    if df is None or df.empty:
        return df, df
    
    split_idx = int(len(df) * train_ratio)
    train_df = df.iloc[:split_idx]
    test_df = df.iloc[split_idx:]
    
    return train_df, test_df


def get_symbol_from_code(code):
    """
    从完整代码中提取股票代码
    """
    if '.' in code:
        return code.split('.')[1]
    return code


def get_exchange_from_code(code):
    """
    从完整代码中提取交易所
    """
    if '.' in code:
        return code.split('.')[0]
    return ''