backend/data/data_service_memory.py

# 数据服务类
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional
import datetime
import logging
import io
from collections import OrderedDict  # 导入 OrderedDict


class DataServiceMemory:
    """
    数据服务类，负责获取和管理股票数据
    """

    def __init__(self, config: Dict[str, Any] = None):
        """
        初始化数据服务

        Args:
            config: 配置信息，包括数据源配置等
        """
        self.config = config or {}
        self.logger = logging.getLogger(__name__)
        # 使用 OrderedDict 实现 LRU 缓存，存储数据
        self.data_cache = OrderedDict()
        self.stock_data = {}
        self.factor_data = {}
        self.backtest_results = {}
        # 定义缓存最大大小
        self.CACHE_MAX_SIZE = 100  # 缓存最大条目数，可根据实际需求调整

    def get_stock_list(self) -> List[Dict[str, Any]]:
        """
        获取股票列表，从API实时获取，不写入本地

        Returns:
            股票列表，包含股票代码、名称等信息
        """
        # 实际项目中应从API实时获取
        # 这里返回示例数据
        return [
            {"code": "000001.SZ", "name": "平安银行", "industry": "银行"},
            {"code": "000002.SZ", "name": "万科A", "industry": "房地产"},
            {"code": "000063.SZ", "name": "中兴通讯", "industry": "通信设备"},
            {"code": "600036.SH", "name": "招商银行", "industry": "银行"},
            {"code": "600519.SH", "name": "贵州茅台", "industry": "白酒"}
        ]

    def get_daily_data(
            self,
            stock_codes: List[str],
            start_date: str,
            end_date: str) -> pd.DataFrame:
        """
        获取日线数据，从API实时获取，并缓存

        Args:
            stock_codes: 股票代码列表
            start_date: 开始日期，格式：YYYY-MM-DD
            end_date: 结束日期，格式：YYYY-MM-DD

        Returns:
            日线数据DataFrame
        """
        cache_key = f"{','.join(stock_codes)}_{start_date}_{end_date}"

        # 检查内存缓存
        if cache_key in self.data_cache:
            self.logger.info(f"从缓存获取日线数据: {cache_key}")
            # 将访问的键移动到 OrderedDict 的末尾，表示最近使用
            self.data_cache.move_to_end(cache_key)
            return self.data_cache[cache_key]

        # 实际项目中应从API实时获取
        # 这里生成示例数据
        self.logger.info(f"从API获取日线数据: {stock_codes}, {start_date} - {end_date}")

        # 生成日期范围
        start = datetime.datetime.strptime(start_date, "%Y-%m-%d")
        end = datetime.datetime.strptime(end_date, "%Y-%m-%d")
        date_range = [
            start +
            datetime.timedelta(
                days=x) for x in range(
                (end -
                 start).days +
                1)]
        date_strings = [d.strftime("%Y-%m-%d") for d in date_range]

        # 生成示例数据
        data = []
        for code in stock_codes:
            base_price = np.random.randint(10, 100)
            for date in date_strings:
                price_change = np.random.normal(0, 1) * base_price * 0.01
                open_price = base_price + price_change
                close_price = open_price + np.random.normal(0, 1) * base_price * 0.005
                high_price = max(open_price, close_price) + \
                    np.random.normal(0, 1) * base_price * 0.003
                low_price = min(open_price, close_price) - \
                    np.random.normal(0, 1) * base_price * 0.003
                volume = np.random.randint(1000000, 10000000)

                data.append({
                    "code": code,
                    "date": date,
                    "open": open_price,
                    "high": high_price,
                    "low": low_price,
                    "close": close_price,
                    "volume": volume,
                    "amount": volume * close_price
                })

                base_price = close_price  # 更新基准价格

        # 创建DataFrame
        df = pd.DataFrame(data)

        # 存入内存缓存
        self.data_cache[cache_key] = df
        # 如果缓存超出最大大小，移除最旧的（即第一个）条目
        if len(self.data_cache) > self.CACHE_MAX_SIZE:
            oldest_key = next(iter(self.data_cache))
            self.data_cache.popitem(last=False)  # 移除最旧的条目
            self.logger.info(f"缓存超出限制，移除最旧条目: {oldest_key}")

        return df

    def get_financial_data(
            self,
            stock_codes: List[str],
            period: str = "latest") -> pd.DataFrame:
        """
        获取财务数据，从API实时获取，并缓存

        Args:
            stock_codes: 股票代码列表
            period: 财报期间，如 'latest', '20220331', '20211231' 等

        Returns:
            财务数据DataFrame
        """
        cache_key = f"financial_{','.join(stock_codes)}_{period}"

        # 检查内存缓存
        if cache_key in self.data_cache:
            self.logger.info(f"从缓存获取财务数据: {cache_key}")
            # 将访问的键移动到 OrderedDict 的末尾，表示最近使用
            self.data_cache.move_to_end(cache_key)
            return self.data_cache[cache_key]

        # 实际项目中应从API实时获取
        # 这里生成示例数据
        self.logger.info(f"从API获取财务数据: {stock_codes}, {period}")

        # 生成示例数据
        data = []
        for code in stock_codes:
            revenue = np.random.randint(1000000, 1000000000)
            net_profit = revenue * np.random.uniform(0.05, 0.2)
            total_assets = revenue * np.random.uniform(1, 5)
            total_equity = total_assets * np.random.uniform(0.3, 0.7)

            data.append({
                "code": code,
                "period": period,
                "revenue": revenue,
                "net_profit": net_profit,
                "total_assets": total_assets,
                "total_equity": total_equity,
                "roe": net_profit / total_equity,
                "roa": net_profit / total_assets,
                "net_margin": net_profit / revenue,
                "debt_ratio": (total_assets - total_equity) / total_assets
            })

        # 创建DataFrame
        df = pd.DataFrame(data)

        # 存入内存缓存
        self.data_cache[cache_key] = df
        # 如果缓存超出最大大小，移除最旧的（即第一个）条目
        if len(self.data_cache) > self.CACHE_MAX_SIZE:
            oldest_key = next(iter(self.data_cache))
            self.data_cache.popitem(last=False)  # 移除最旧的条目
            self.logger.info(f"缓存超出限制，移除最旧条目: {oldest_key}")

        return df

    def save_backtest_result(self, strategy_id: str, result_data: Dict[str, Any]):
        """
        保存回测结果

        Args:
            strategy_id: 策略ID
            result_data: 回测结果数据
        """
        self.backtest_results[strategy_id] = result_data
        self.logger.info(f"回测结果已保存: {strategy_id}")

    def get_backtest_result(self, strategy_id: str) -> Optional[Dict[str, Any]]:
        """
        从缓存获取回测结果

        Args:
            strategy_id: 策略ID

        Returns:
            回测结果数据，如果不存在则返回None
        """
        return self.backtest_results.get(strategy_id)

    def clear_cache(self):
        """
        清除内存缓存
        """
        self.data_cache.clear()
        self.logger.info("缓存已清除")

    def export_data(self, data: pd.DataFrame, format: str = "csv") -> bytes:
        """
        导出数据为字节流

        Args:
            data: 要导出的数据
            format: 导出格式，支持 'csv', 'excel', 'json'

        Returns:
            字节数据
        """
        buffer = io.BytesIO()

        if format.lower() == "csv":
            data.to_csv(buffer, index=False, encoding="utf-8")
        elif format.lower() == "excel":
            data.to_excel(buffer, index=False, engine="openpyxl")
        elif format.lower() == "json":
            buffer.write(data.to_json(orient="records").encode("utf-8"))
        else:
            raise ValueError(f"不支持的导出格式: {format}")

        buffer.seek(0)
        return buffer.getvalue()