fix: keep baostock session alive across bulk downloads
baostock drops a session after a few hundred queries; every later query then returned 用户未登录 and the symbol failed. download_daily_batch now refreshes the session every relogin_every symbols and re-logs in + retries once on a detected session loss. akshare fallback now defaults off in the batch path since it is slow/unreliable on this network and the re-login keeps baostock as the fast primary. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
+56
-22
@@ -13,6 +13,10 @@ _BAOSTOCK_FIELDS = "date,open,high,low,close,volume,amount"
|
||||
_OHLCV = ["open", "high", "low", "close", "volume", "amount"]
|
||||
|
||||
|
||||
class _SessionLost(Exception):
|
||||
"""baostock reported the session was dropped (``用户未登录``)."""
|
||||
|
||||
|
||||
def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
|
||||
"""Download daily bars from akshare. Returns DataFrame with OHLCV columns."""
|
||||
try:
|
||||
@@ -124,50 +128,80 @@ def download_daily_batch(
|
||||
start: str,
|
||||
end: str,
|
||||
adjust: str = "qfq",
|
||||
akshare_fallback: bool = True,
|
||||
akshare_fallback: bool = False,
|
||||
relogin_every: int = 200,
|
||||
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
|
||||
"""Download many symbols under a single baostock session.
|
||||
"""Download many symbols, keeping a baostock session alive across the run.
|
||||
|
||||
Logging into baostock once per call (instead of per symbol) is the dominant
|
||||
speed-up when fetching thousands of symbols. Yields ``(symbol, df)`` as each
|
||||
symbol completes so callers can stream results to disk; ``df`` is ``None``
|
||||
when both sources fail. Each ``df`` has the same 8 columns as
|
||||
:func:`download_daily`.
|
||||
Logging in once (instead of per symbol) is the dominant speed-up for
|
||||
thousands of symbols, but baostock drops a session after a while
|
||||
(subsequent queries return ``用户未登录``). So we refresh the session every
|
||||
``relogin_every`` symbols and also re-login + retry once whenever a query
|
||||
reports the session is gone. Yields ``(symbol, df)`` as each symbol
|
||||
completes; ``df`` is ``None`` when no data is available. Each ``df`` has the
|
||||
same 8 columns as :func:`download_daily`.
|
||||
|
||||
Args:
|
||||
symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
|
||||
start, end: ``YYYY-MM-DD`` bounds.
|
||||
adjust: ``qfq`` / ``hfq`` / ``''``.
|
||||
akshare_fallback: Retry a failed symbol through akshare before yielding
|
||||
``None``.
|
||||
akshare_fallback: Retry a failed symbol through akshare. Off by default
|
||||
because akshare is unreliable on the deployment network and each
|
||||
failed attempt is slow; baostock + re-login is the fast path.
|
||||
relogin_every: Proactively refresh the baostock session every N symbols.
|
||||
"""
|
||||
flag = _BAOSTOCK_ADJUST.get(adjust, "2")
|
||||
|
||||
def _relogin() -> None:
|
||||
try:
|
||||
bs.logout()
|
||||
except Exception:
|
||||
pass
|
||||
bs.login()
|
||||
try:
|
||||
for symbol in symbols:
|
||||
df: Optional[pd.DataFrame] = None
|
||||
try:
|
||||
|
||||
def _fetch(symbol: str) -> Optional[pd.DataFrame]:
|
||||
"""One baostock query; returns df, or None (no data), or raises _SessionLost."""
|
||||
code = f"{symbol[:2]}.{symbol[2:]}"
|
||||
rs = bs.query_history_k_data_plus(
|
||||
code=code, fields=_BAOSTOCK_FIELDS,
|
||||
start_date=start, end_date=end,
|
||||
frequency="d", adjustflag=flag,
|
||||
start_date=start, end_date=end, frequency="d", adjustflag=flag,
|
||||
)
|
||||
if rs.error_code == "0":
|
||||
if rs.error_code != "0":
|
||||
if "未登录" in (rs.error_msg or ""):
|
||||
raise _SessionLost(rs.error_msg)
|
||||
logger.warning("baostock error for %s: %s", symbol, rs.error_msg)
|
||||
return None
|
||||
rows = []
|
||||
while rs.next():
|
||||
rows.append(rs.get_row_data())
|
||||
if rows:
|
||||
if not rows:
|
||||
return None
|
||||
df = pd.DataFrame(rows, columns=["date", *_OHLCV])
|
||||
# Suspended-trading days come back as empty strings;
|
||||
# coerce to NaN rather than crashing the whole symbol.
|
||||
# Suspended-trading days come back as empty strings; coerce to NaN
|
||||
# rather than crashing the whole symbol.
|
||||
df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce")
|
||||
df["symbol"] = symbol
|
||||
df = df[["symbol", "date", *_OHLCV]]
|
||||
else:
|
||||
logger.warning("baostock error for %s: %s", symbol, rs.error_msg)
|
||||
return df[["symbol", "date", *_OHLCV]]
|
||||
|
||||
bs.login()
|
||||
try:
|
||||
for i, symbol in enumerate(symbols):
|
||||
if i and relogin_every and i % relogin_every == 0:
|
||||
_relogin()
|
||||
|
||||
df: Optional[pd.DataFrame] = None
|
||||
for attempt in (1, 2):
|
||||
try:
|
||||
df = _fetch(symbol)
|
||||
break
|
||||
except _SessionLost:
|
||||
if attempt == 1:
|
||||
_relogin() # session dropped — refresh and retry once
|
||||
continue
|
||||
logger.warning("baostock session lost for %s after relogin", symbol)
|
||||
except Exception as e:
|
||||
logger.warning("baostock download failed for %s: %s", symbol, e)
|
||||
break
|
||||
|
||||
if (df is None or df.empty) and akshare_fallback:
|
||||
df = _download_akshare(symbol, start, end, adjust)
|
||||
|
||||
Reference in New Issue
Block a user