feat(data): DataFeed 多源 fallback + BaoStock 超时修复

sanguo_data/datafeed.py

@@ -0,0 +1,228 @@
+# sanguo_data/datafeed.py
+import pandas as pd
+import urllib.request
+import json
+import time
+import logging
+from datetime import datetime, timedelta
+from multiprocessing import Process, Queue
+from typing import Optional
+from sanguo_data.config import DataConfig
+
+logger = logging.getLogger(__name__)
+
+def fetch_with_fallback(symbol, start, end, sources: list[str]) -> pd.DataFrame:
+    fetchers = {
+        "eastmoney": _fetch_eastmoney,
+        "baostock": lambda s, st, e: _fetch_baostock_with_timeout(s, st, e, timeout=30),
+        "tencent": _fetch_tencent,
+    }
+    last_err = None
+    for name in sources:
+        try:
+            df = fetchers[name](symbol, start, end)
+            if df is not None and len(df) > 0:
+                return df
+        except Exception as e:
+            last_err = e
+            continue
+    raise RuntimeError(f"all sources failed: {last_err}")
+
+def fetch_daily(symbol, start, end, cfg: DataConfig) -> pd.DataFrame:
+    sources = [s["name"] for s in cfg.data_sources.get("daily", []) if s.get("enabled", True)]
+    return fetch_with_fallback(symbol, start, end, sources)
+
+# Worker function for multiprocessing (must be at module level to be picklable)
+def _baostock_worker(symbol, start, end, result_queue):
+    try:
+        result = _fetch_baostock_raw(symbol, start, end)
+        result_queue.put(result)
+    except Exception as e:
+        result_queue.put(e)
+
+# Test helper for timeout testing (simulates hanging BaoStock call)
+def _hanging_worker_for_test(symbol, start, end, result_queue):
+    """Test helper: simulates a hanging BaoStock call (60s sleep)"""
+    import time as _time
+    _time.sleep(60)
+    result_queue.put(pd.DataFrame({"date": ["2026-01-01"], "open": [10.0]}))
+
+def _fetch_baostock_with_timeout(symbol, start, end, timeout=30):
+    """子进程隔离 BaoStock（修复 v1 无超时卡死坑）"""
+    result_queue = Queue()
+    p = Process(target=_baostock_worker, args=(symbol, start, end, result_queue))
+    p.start()
+    p.join(timeout)
+    if p.is_alive():
+        p.terminate()
+        p.join()
+        raise TimeoutError(f"baostock timeout after {timeout}s")
+
+    res = result_queue.get()
+    if isinstance(res, Exception):
+        raise res
+    return res
+
+def _get_em_secid(code: str) -> str:
+    if code.startswith(("60", "68", "51")):
+        return f"1.{code}"
+    return f"0.{code}"
+
+def _parse_em_klines(klines: list) -> Optional[pd.DataFrame]:
+    """解析东方财富K线数据（日线和15min通用）"""
+    if not klines:
+        return None
+    rows = []
+    for line in klines:
+        parts = line.split(",")
+        if len(parts) < 7:
+            continue
+        rows.append({
+            "date": parts[0],
+            "open": float(parts[1]),
+            "close": float(parts[2]),
+            "high": float(parts[3]),
+            "low": float(parts[4]),
+            "volume": float(parts[5]),
+            "amount": float(parts[6]),
+        })
+    if not rows:
+        return None
+    return pd.DataFrame(rows)
+
+def _fetch_baostock_raw(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
+    """BaoStock日线：全量历史，无反爬，amount真实，T+1延迟
+
+    Copied from v1 data_platform/daily_all_update.py:fetch_baostock_daily (lines 242-270)
+    """
+    try:
+        import baostock as bs
+    except ImportError:
+        return None
+
+    # 转换代码格式：600000 -> sh.600000
+    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
+    if code.startswith(("60", "68", "51")):
+        bs_code = f"sh.{code}"
+    else:
+        bs_code = f"sz.{code}"
+
+    try:
+        rs = bs.query_history_k_data_plus(
+            bs_code,
+            "date,open,high,low,close,volume,amount",
+            start_date=start_date.replace("-", ""),
+            end_date=end_date.replace("-", ""),
+            frequency="d",
+            adjustflag="2",
+        )
+        rows = []
+        while (rs.error_code == "0") and rs.next():
+            rows.append(rs.get_row_data())
+        if not rows:
+            return None
+        df = pd.DataFrame(rows, columns=["date", "open", "high", "low", "close", "volume", "amount"])
+        for c in ["open", "high", "low", "close", "volume", "amount"]:
+            df[c] = pd.to_numeric(df[c], errors="coerce")
+        df = df.dropna(subset=["close"])
+        if df.empty:
+            return None
+        return df
+    except Exception as e:
+        logger.debug("BaoStock日线失败 %s: %s", symbol, e)
+        return None
+
+def _fetch_eastmoney(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
+    """东方财富日线：当天实时，amount真实，4s限频
+
+    Copied from v1 data_platform/daily_all_update.py:fetch_eastmoney_daily (lines 339-374)
+    """
+    try:
+        import requests as _requests
+    except ImportError:
+        return None
+
+    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
+    secid = _get_em_secid(code)
+    ts = str(int(time.time() * 1000))
+    url = (
+        f"https://push2his.eastmoney.com/api/qt/stock/kline/get?"
+        f"secid={secid}&klt=101&fqt=1&"
+        f"beg={start_date.replace('-', '')}&end={end_date.replace('-', '')}&"
+        f"fields1=f1,f2,f3,f4,f5,f6,f7,f8&"
+        f"fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&"
+        f"ut=b2884a393a59ad64002292a3e90d46a5&lmt=10000&"
+        f"cb=jQuery_em_{ts}&_={ts}"
+    )
+
+    headers_em = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
+        "Referer": "https://quote.eastmoney.com/",
+        "Accept": "*/*",
+        "Accept-Language": "zh-CN,zh;q=0.9",
+    }
+
+    session = _requests.Session()
+    session.trust_env = False
+    try:
+        r = session.get(url, headers=headers_em, timeout=15, verify=False)
+        if r.status_code != 200:
+            return None
+        text = r.text
+        data = json.loads(text[text.index("(") + 1:text.rindex(")")])
+        if data.get("rc") != 0:
+            return None
+        klines = data.get("data", {}).get("klines", [])
+        df = _parse_em_klines(klines)
+        if df is None:
+            return None
+        df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d")
+        mask = (df["date"] >= start_date) & (df["date"] <= end_date)
+        result = df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]]
+        return result if not result.empty else None
+    except Exception as e:
+        logger.debug("东方财富日线失败 %s: %s", symbol, e)
+        return None
+
+def _fetch_tencent(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
+    """腾讯日线：amount有时为0
+
+    Copied from v1 data_platform/fallback.py:_fetch_tencent_daily (lines 66-104)
+    """
+    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
+    if code.startswith(("6", "5", "1")):
+        prefix = "sh"
+    else:
+        prefix = "sz"
+    tq_symbol = f"{prefix}{code}"
+
+    days = (datetime.strptime(end_date, "%Y-%m-%d") - datetime.strptime(start_date, "%Y-%m-%d")).days + 10
+    url = f"https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={tq_symbol},day,{start_date},,{days},"
+
+    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
+    try:
+        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
+        with opener.open(req, timeout=10) as r:
+            resp = json.loads(r.read())
+        d = resp.get("data")
+        if not isinstance(d, dict):
+            return None
+        klines = d.get(tq_symbol, {}).get("day", [])
+        if not klines:
+            return None
+        df = pd.DataFrame(klines)
+        ncols = len(df.columns)
+        if ncols >= 7:
+            df.columns = ["date", "open", "close", "high", "low", "volume", "amount"][:ncols]
+        else:
+            df.columns = ["date", "open", "close", "high", "low", "volume"][:ncols]
+        if "amount" not in df.columns:
+            df["amount"] = 0.0
+        for c in ["open", "close", "high", "low", "volume", "amount"]:
+            df[c] = pd.to_numeric(df[c], errors="coerce").fillna(0)
+        df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d")
+        mask = (df["date"] >= start_date) & (df["date"] <= end_date)
+        return df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]]
+    except Exception as e:
+        logger.debug("腾讯日线失败 %s: %s", symbol, e)
+        return None

tests/data/test_datafeed.py

@@ -0,0 +1,49 @@
+# tests/data/test_datafeed.py
+import time
+import pandas as pd
+import pytest
+from unittest.mock import patch
+from sanguo_data.datafeed import fetch_with_fallback, _fetch_baostock_with_timeout
+
+def test_fetch_with_fallback_uses_second_when_first_fails():
+    df_good = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0], "high": [11.0], "low": [9.0], "close": [10.5], "volume": [1000], "amount": [10000]})
+    with patch("sanguo_data.datafeed._fetch_eastmoney", side_effect=Exception("limit")), \
+         patch("sanguo_data.datafeed._fetch_baostock_with_timeout", return_value=df_good):
+        out = fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
+    assert len(out) == 1
+    assert out["date"].iloc[0] == "2026-01-01"
+
+def test_baostock_timeout_does_not_hang():
+    """v1 卡死坑修复验证：超时必须返回，不能无限挂起"""
+    import sanguo_data.datafeed as df_module
+
+    # Save original worker
+    original_worker = df_module._baostock_worker
+
+    try:
+        # Replace with hanging worker (module-level function can be pickled)
+        df_module._baostock_worker = df_module._hanging_worker_for_test
+
+        start = time.time()
+        with pytest.raises(TimeoutError):
+            df_module._fetch_baostock_with_timeout("600000", "2026-01-01", "2026-01-02", timeout=2)
+        elapsed = time.time() - start
+        assert elapsed < 5, f"Timeout test took {elapsed:.2f}s, expected <5s"
+    finally:
+        # Restore original worker
+        df_module._baostock_worker = original_worker
+
+def test_fetch_with_fallback_all_sources_fail():
+    """所有源都失败时应该抛出异常"""
+    with patch("sanguo_data.datafeed._fetch_eastmoney", side_effect=Exception("em failed")), \
+         patch("sanguo_data.datafeed._fetch_baostock_with_timeout", side_effect=Exception("bs failed")):
+        with pytest.raises(RuntimeError, match="all sources failed"):
+            fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
+
+def test_fetch_with_fallback_first_succeeds():
+    """第一个源成功时直接返回"""
+    df_good = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0], "high": [11.0], "low": [9.0], "close": [10.5], "volume": [1000], "amount": [10000]})
+    with patch("sanguo_data.datafeed._fetch_eastmoney", return_value=df_good):
+        out = fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
+    assert len(out) == 1
+    assert out["date"].iloc[0] == "2026-01-01"