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"}
_BAOSTOCK_FIELDS = "date,open,high,low,close,volume,amount"
_OHLCV = ["open", "high", "low", "close", "volume", "amount"]


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 = True,
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
    """Download many symbols under a single baostock session.

    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`.

    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``.
    """
    flag = _BAOSTOCK_ADJUST.get(adjust, "2")
    bs.login()
    try:
        for symbol in symbols:
            df: Optional[pd.DataFrame] = None
            try:
                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,
                )
                if rs.error_code == "0":
                    rows = []
                    while rs.next():
                        rows.append(rs.get_row_data())
                    if rows:
                        df = pd.DataFrame(rows, columns=["date", *_OHLCV])
                        # 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)
            except Exception as e:
                logger.warning("baostock download failed for %s: %s", symbol, e)

            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