""" DataManager — 统一数据管理层。 策略/模型层通过 DataManager 获取数据,不直接访问 AkShare/Tushare 或数据库。 优先从 DB 读取,缺失时依次尝试 AkShare → Tushare 拉取并入库。 """ import logging import time import pandas as pd from config.settings import DEFAULT_START_DATE, DEFAULT_END_DATE from data.sources.akshare_source import AkShareSource from data.sources.tushare_source import TushareSource from database import dao from database.models import create_all_tables logger = logging.getLogger(__name__) def _rollback_days(yyyymmdd: str, days: int) -> str: """把 'YYYYMMDD' 往前回退 days 个自然日,返回 'YYYYMMDD'。""" from datetime import datetime, timedelta dt = datetime.strptime(yyyymmdd, "%Y%m%d") - timedelta(days=days) return dt.strftime("%Y%m%d") def _safe_save_daily(df: pd.DataFrame, skip_exc: bool = True) -> None: """把拉取到的日线入库;失败时记日志(可选抛错)。""" cols = [c for c in dao._DAILY_COLS if c in df.columns] try: dao.save_daily(df[cols]) except Exception as e: if skip_exc: logger.exception("[DataManager] 日线入库失败(数据已获取,未持久化):%s", e) else: raise class DataManager: """统一数据管理。双数据源:AkShare(主)+ Tushare(备)。""" def __init__(self): self._ak_source: AkShareSource | None = None self._ts_source: TushareSource | None = None @property def ak(self) -> AkShareSource: if self._ak_source is None: self._ak_source = AkShareSource() return self._ak_source @property def ts(self) -> TushareSource: if self._ts_source is None: self._ts_source = TushareSource() return self._ts_source # ── 初始化 ──────────────────────────────────────────── def init_db(self) -> None: create_all_tables() # ── 内部:双源 try ──────────────────────────────────── def _try_fetch(self, method_name: str, *args, **kwargs): """ 依次尝试 AkShare → Tushare 调用同一方法名。 method_name: 'fetch_daily' | 'fetch_stock_list' | 'fetch_financial' 返回: (result_df, source_name) 或 (empty_df, None) 如果 ts_code 是指数代码,自动路由到 fetch_index_daily。 """ from data.sources.akshare_source import is_index_code # 指数自动路由 if method_name == "fetch_daily" and args: ts_code = args[0] if is_index_code(ts_code): method_name = "fetch_index_daily" for label, source in [("Tushare", self.ts), ("AkShare", self.ak)]: try: if label == "Tushare" and not source.available: continue fn = getattr(source, method_name) df = fn(*args, **kwargs) if df is not None and not df.empty: return df, label except Exception as e: logger.warning(" [%s] %s 失败: %s", label, method_name, e) return pd.DataFrame(), None # ── 股票列表 ────────────────────────────────────────── def get_stock_list(self, force_refresh: bool = False) -> pd.DataFrame: if not force_refresh: df = dao.query_stock_list() if not df.empty: logger.info("[DataManager] 从 DB 读取股票列表: %s 只", len(df)) return df logger.info("[DataManager] 拉取股票列表 (AkShare → Tushare)...") df, src = self._try_fetch("fetch_stock_list") if df.empty: logger.warning("[DataManager] 所有数据源均无法获取股票列表") return pd.DataFrame() logger.info("[DataManager] 股票列表已入库 (%s): %s 只", src, len(df)) dao.save_stock_list(df) time.sleep(2) return df # ── 日线数据 ────────────────────────────────────────── def get_daily( self, ts_code: str, start: str | None = None, end: str | None = None, force_refresh: bool = False, ) -> pd.DataFrame: start = start or DEFAULT_START_DATE end = end or DEFAULT_END_DATE or time.strftime("%Y%m%d") if not force_refresh: df = dao.query_daily(ts_code, start, end) if not df.empty: # 检查 DB 窗口尾部是否明显落后于请求的 end(>5 个自然日)。 # 若是,说明存在数据缺口,触发一次增量补拉,而不是把陈旧数据当完整返回。 latest = str(df["trade_date"].astype(str).max()) if end >= "20200101" and latest < _rollback_days(end, 7): logger.warning( "%s DB 窗口尾部 %s 明显落后于请求日 %s,触发补拉", ts_code, latest, end) refreshed, src = self._try_fetch("fetch_daily", ts_code, latest, end) if not refreshed.empty and src: _safe_save_daily(refreshed) return df # DB 未命中 → 从数据源拉取 df, src = self._try_fetch("fetch_daily", ts_code, start, end) if df.empty: logger.warning("%s 日线获取失败 (AkShare+Tushare 均不可用)", ts_code) return pd.DataFrame() # 保存到 DB(失败仅记日志,不阻断已获取到的数据返回) _safe_save_daily(df) return df def sync_daily(self, ts_code: str) -> int: latest = dao.get_latest_trade_date(ts_code) today = time.strftime("%Y%m%d") if latest and latest >= today: logger.info("[DataManager] %s 数据已是最新 (%s)", ts_code, latest) return 0 start = latest or DEFAULT_START_DATE df, src = self._try_fetch("fetch_daily", ts_code, start, today) if df.empty: logger.warning("[DataManager] %s sync_daily 失败 (AkShare+Tushare 均不可用)", ts_code) return 0 # 入库失败要能反映到返回值,否则调用方会误以为已持久化 cols = [c for c in dao._DAILY_COLS if c in df.columns] try: dao.save_daily(df[cols]) except Exception as e: logger.exception("[DataManager] %s 日线入库失败:%s", ts_code, e) raise logger.info("[DataManager] %s 同步 %s 条日线 (%s)", ts_code, len(df), src) return len(df) def sync_all_daily(self) -> int: stock_list = self.get_stock_list() total = 0 for i, ts_code in enumerate(stock_list.index): try: total += self.sync_daily(ts_code) if (i + 1) % 50 == 0: logger.info("[DataManager] 进度: %s/%s", i + 1, len(stock_list)) time.sleep(1) except Exception as e: logger.warning("[DataManager] %s 同步失败: %s", ts_code, e) logger.info("[DataManager] 全量同步完成,新增 %s 条", total) return total # ── 财务数据 ────────────────────────────────────────── def get_financial(self, ts_code: str) -> pd.DataFrame: df = dao.query_financial(ts_code) if not df.empty: return df df, src = self._try_fetch("fetch_financial", ts_code) if not df.empty: try: cols = [c for c in dao._FINA_COLS if c in df.columns] dao.save_financial(df[cols]) except Exception as e: logger.exception("财务数据入库失败: %s", e) return df