chinese-equity-quant/data/downloader.py

"""Unified data downloader: baostock primary, akshare fallback."""
import logging
from typing import Iterable, Iterator, Optional, Tuple
import pandas as pd
import akshare as ak
import baostock as bs

logger = logging.getLogger(__name__)

# Map the adjust argument to baostock's adjustflag codes.
_BAOSTOCK_ADJUST = {"qfq": "2", "hfq": "1", "": "3", "none": "3"}

# Richer field set requested by the batch downloader. On top of OHLCV+amount we
# pull baostock's preclose, turnover rate, daily % change, trade/ST status, and
# the four valuation ratios, then derive a daily VWAP (amount / volume).
_BATCH_FIELDS = (
    "date,open,high,low,close,preclose,volume,amount,turn,pctChg,"
    "tradestatus,isST,peTTM,pbMRQ,psTTM,pcfNcfTTM"
)
# Every batch field except ``date`` is numeric (flags included: 0/1 strings).
_BATCH_NUMERIC = [
    "open", "high", "low", "close", "preclose", "volume", "amount",
    "turn", "pctChg", "tradestatus", "isST",
    "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
]
# Output column order; ``vwap`` is derived (inserted right after ``amount``).
_BATCH_COLUMNS = [
    "symbol", "date",
    "open", "high", "low", "close", "preclose", "volume", "amount", "vwap",
    "turn", "pctChg", "tradestatus", "isST",
    "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
]

# Raw Baostock minute bars. The ``time`` field is usually compact
# YYYYMMDDHHMMSSmmm; parsing below also tolerates HH:MM:SS strings in tests.
_MINUTE_FIELDS = "date,time,code,open,high,low,close,volume,amount,adjustflag"
_MINUTE_NUMERIC = ["open", "high", "low", "close", "volume", "amount"]
_MINUTE_COLUMNS = [
    "symbol", "datetime", "date", "time", "frequency",
    "open", "high", "low", "close", "volume", "amount", "vwap", "adjustflag",
]
_MINUTE_FREQUENCIES = {"5", "15", "30", "60"}


class _SessionLost(Exception):
    """baostock reported the session was dropped (``用户未登录``)."""


def _normalize_minute_frequency(frequency: str | int) -> tuple[str, str]:
    """Return Baostock frequency and partition label for a minute interval."""
    raw = str(frequency).strip().lower()
    if raw.endswith("m"):
        raw = raw[:-1]
    if raw not in _MINUTE_FREQUENCIES:
        raise ValueError(
            f"Unsupported minute frequency '{frequency}'. "
            f"Expected one of {sorted(_MINUTE_FREQUENCIES)} minutes."
        )
    return raw, f"{raw}m"


def _parse_minute_datetime(date: pd.Series, time: pd.Series) -> pd.Series:
    """Parse Baostock minute timestamps into pandas datetimes."""
    date_dt = pd.to_datetime(date, errors="coerce")
    date_compact = date_dt.dt.strftime("%Y%m%d")
    time_text = time.astype(str).str.strip()
    time_digits = time_text.str.replace(r"\D", "", regex=True)

    full_digits = time_digits.str.slice(0, 14)
    from_full = pd.to_datetime(full_digits, format="%Y%m%d%H%M%S", errors="coerce")

    from_short = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]")
    short_time = time_digits.str.len().between(1, 6)
    if short_time.any():
        short_digits = (
            time_digits.loc[short_time]
            .str.pad(6, side="right", fillchar="0")
            .str.slice(0, 6)
        )
        from_short.loc[short_time] = pd.to_datetime(
            date_compact.loc[short_time] + short_digits,
            format="%Y%m%d%H%M%S",
            errors="coerce",
        )

    from_text = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]")
    text_time = time_text.str.contains(":", regex=False)
    if text_time.any():
        from_text.loc[text_time] = pd.to_datetime(
            date.astype(str).loc[text_time] + " " + time_text.loc[text_time],
            errors="coerce",
        )
    return from_full.fillna(from_short).fillna(from_text)


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:
        # symbol format: 'sh600000' in akshare stock_zh_a_hist expects raw code like '600000'
        # strip exchange prefix for akshare
        raw = symbol.replace("sh", "").replace("sz", "")
        df = ak.stock_zh_a_hist(
            symbol=raw,
            period="daily",
            start_date=start,
            end_date=end,
            adjust=adjust,
        )
        if df is None or df.empty:
            return None
        # Standardize columns
        col_map = {
            "日期": "date",
            "开盘": "open",
            "最高": "high",
            "最低": "low",
            "收盘": "close",
            "成交量": "volume",
            "成交额": "amount",
        }
        df = df.rename(columns=col_map)
        df["symbol"] = symbol
        return df[["symbol", "date", "open", "high", "low", "close", "volume", "amount"]]
    except Exception as e:
        logger.warning(f"akshare download failed for {symbol}: {e}")
        return None


def _download_baostock(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
    """Download daily bars from baostock (primary source)."""
    try:
        bs.login()
        # baostock format: sh.600000
        code = f"{symbol[:2]}.{symbol[2:]}"
        rs = bs.query_history_k_data_plus(
            code=code,
            fields="date,open,high,low,close,volume,amount",
            start_date=start,
            end_date=end,
            frequency="d",
            adjustflag=_BAOSTOCK_ADJUST.get(adjust, "2"),
        )
        if rs.error_code != "0":
            logger.warning(f"baostock error for {symbol}: {rs.error_msg}")
            return None
        data_list = []
        while rs.next():
            data_list.append(rs.get_row_data())
        if not data_list:
            return None
        df = pd.DataFrame(data_list, columns=["date", "open", "high", "low", "close", "volume", "amount"])
        df[["open", "high", "low", "close", "volume", "amount"]] = df[
            ["open", "high", "low", "close", "volume", "amount"]
        ].apply(pd.to_numeric, errors="coerce")
        df["symbol"] = symbol
        return df[["symbol", "date", "open", "high", "low", "close", "volume", "amount"]]
    except Exception as e:
        logger.warning(f"baostock download failed for {symbol}: {e}")
        return None
    finally:
        try:
            bs.logout()
        except Exception:
            pass


def download_daily(
    symbol: str,
    start: str,
    end: str,
    adjust: str = "qfq",
    source: str = "auto",
) -> pd.DataFrame:
    """
    Download daily OHLCV data. Tries baostock first, falls back to akshare.

    Args:
        symbol: Stock symbol like 'sh600000' or 'sz000001'
        start: Start date 'YYYY-MM-DD'
        end: End date 'YYYY-MM-DD'
        adjust: 'qfq' (forward-adjusted), 'hfq' (backward), '' (none)
        source: 'auto' (baostock then akshare fallback), 'baostock' only,
            or 'akshare' only

    Returns:
        DataFrame with columns: symbol, date, open, high, low, close, volume, amount
    """
    df = None
    if source in ("baostock", "auto"):
        df = _download_baostock(symbol, start, end, adjust)
    if df is None and source in ("akshare", "auto"):
        df = _download_akshare(symbol, start, end, adjust)

    if df is None or df.empty:
        raise RuntimeError(f"Failed to download data for {symbol} from {start} to {end}")

    df["date"] = pd.to_datetime(df["date"])
    df = df.sort_values("date").reset_index(drop=True)
    return df


def download_daily_batch(
    symbols: Iterable[str],
    start: str,
    end: str,
    adjust: str = "qfq",
    akshare_fallback: bool = False,
    relogin_every: int = 200,
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
    """Download many symbols, keeping a baostock session alive across the run.

    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. 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()

    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=_BATCH_FIELDS,
            start_date=start, end_date=end, frequency="d", adjustflag=flag,
        )
        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 not rows:
            return None
        df = pd.DataFrame(rows, columns=["date", *_BATCH_NUMERIC])
        # Suspended-trading days come back as empty strings; coerce to NaN
        # rather than crashing the whole symbol.
        df[_BATCH_NUMERIC] = df[_BATCH_NUMERIC].apply(pd.to_numeric, errors="coerce")
        # Daily VWAP = turnover (yuan) / shares; NaN when no volume (suspended).
        df["vwap"] = (df["amount"] / df["volume"]).where(df["volume"] > 0)
        df["symbol"] = symbol
        return df[_BATCH_COLUMNS]

    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)

            if df is not None and not df.empty:
                df["date"] = pd.to_datetime(df["date"])
                df = df.sort_values("date").reset_index(drop=True)
                yield symbol, df
            else:
                yield symbol, None
    finally:
        try:
            bs.logout()
        except Exception:
            pass


def download_minute_batch(
    symbols: Iterable[str],
    start: str,
    end: str,
    frequency: str | int = 5,
    relogin_every: int = 200,
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
    """Download raw Baostock minute bars for many symbols.

    Minute bars are intentionally unadjusted (`adjustflag='3'`) because the
    output is raw intraday market data for downstream feature aggregation, not a
    tradable daily price series.

    Args:
        symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
        start, end: ``YYYY-MM-DD`` bounds.
        frequency: Baostock minute frequency. ``5``/``"5"``/``"5m"`` all mean
            5-minute bars.
        relogin_every: Proactively refresh the baostock session every N symbols.

    Yields:
        ``(symbol, df)`` where ``df`` has raw minute bars or ``None`` when no
        data is available.
    """
    query_frequency, frequency_label = _normalize_minute_frequency(frequency)
    adjustflag = _BAOSTOCK_ADJUST["none"]

    def _relogin() -> None:
        try:
            bs.logout()
        except Exception:
            pass
        bs.login()

    def _fetch(symbol: str) -> Optional[pd.DataFrame]:
        """One Baostock minute query; returns df, None, or raises _SessionLost."""
        code = f"{symbol[:2]}.{symbol[2:]}"
        rs = bs.query_history_k_data_plus(
            code=code,
            fields=_MINUTE_FIELDS,
            start_date=start,
            end_date=end,
            frequency=query_frequency,
            adjustflag=adjustflag,
        )
        if rs.error_code != "0":
            if "未登录" in (rs.error_msg or ""):
                raise _SessionLost(rs.error_msg)
            logger.warning("baostock minute error for %s: %s", symbol, rs.error_msg)
            return None
        rows = []
        while rs.next():
            rows.append(rs.get_row_data())
        if not rows:
            return None

        df = pd.DataFrame(rows, columns=_MINUTE_FIELDS.split(","))
        df[_MINUTE_NUMERIC] = df[_MINUTE_NUMERIC].apply(pd.to_numeric, errors="coerce")
        df["datetime"] = _parse_minute_datetime(df["date"], df["time"])
        bad_timestamps = df["datetime"].isna()
        if bad_timestamps.any():
            raise ValueError(
                f"Could not parse {int(bad_timestamps.sum())} minute timestamp(s)"
            )
        df["date"] = pd.to_datetime(df["date"], errors="coerce")
        df["time"] = df["datetime"].dt.strftime("%H:%M:%S")
        df["frequency"] = frequency_label
        df["vwap"] = (df["amount"] / df["volume"]).where(df["volume"] > 0)
        df["symbol"] = symbol
        return df[_MINUTE_COLUMNS].sort_values("datetime").reset_index(drop=True)

    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()
                        continue
                    logger.warning("baostock minute session lost for %s after relogin", symbol)
                except Exception as e:
                    logger.warning("baostock minute download failed for %s: %s", symbol, e)
                    break

            if df is not None and not df.empty:
                yield symbol, df
            else:
                yield symbol, None
    finally:
        try:
            bs.logout()
        except Exception:
            pass