From 3015d6fa6bc5bf4bea98fd4c396c314befb10797 Mon Sep 17 00:00:00 2001 From: claude_dev Date: Sun, 5 Jul 2026 12:41:00 +0800 Subject: [PATCH] =?UTF-8?q?feat(data):=20DataFeed=20=E5=A4=9A=E6=BA=90=20f?= =?UTF-8?q?allback=20+=20BaoStock=20=E8=B6=85=E6=97=B6=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- sanguo_data/datafeed.py | 228 ++++++++++++++++++++++++++++++++++++ tests/data/test_datafeed.py | 49 ++++++++ 2 files changed, 277 insertions(+) create mode 100644 sanguo_data/datafeed.py create mode 100644 tests/data/test_datafeed.py diff --git a/sanguo_data/datafeed.py b/sanguo_data/datafeed.py new file mode 100644 index 0000000..a3f833b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/data/test_datafeed.py b/tests/data/test_datafeed.py new file mode 100644 index 0000000..fc69d73 --- /dev/null +++ b/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"