# 数据服务类 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()