mirror of
https://ghfast.top/https://github.com/aeroxw/easy_tdx_max.git
synced 2026-09-12 16:54:20 +08:00
feat: /bars 自动兜底数据源 baostock — TDX 全部路径失败时最后一级回退（仅日线及以上），响应带 source 字段标注数据来源
This commit is contained in:
GitHub
@@ -21,6 +21,9 @@ science = ["scipy>=1.10,<1.16"]
|
||||
web = ["fastapi>=0.110,<1", "uvicorn[standard]>=0.29"]
|
||||
# 本地 K 线数据仓库（easy-tdx warehouse ...）：DuckDB 单文件列存
|
||||
warehouse = ["duckdb>=1.0"]
|
||||
# baostock 自动兜底数据源（TDX 全部路径失败时的最后一级回退，仅日线及以上）。
|
||||
# 可选安装：不装则该回退环自动关闭，核心功能不受影响。
|
||||
baostock = ["baostock>=0.9"]
|
||||
# 打包成桌面 EXE 用：系统托盘（pystray）+ 图标生成（Pillow）。
|
||||
# 仅 PyInstaller 打包态需要，开发态 ``pip install -e .[web]`` 不强制装。
|
||||
packaging = ["pystray>=0.19", "Pillow>=10.0"]
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""备选数据源（自动兜底）。"""
|
||||
@@ -0,0 +1,171 @@
|
||||
"""baostock 自动兜底数据源（仅日线及以上，TDX 全部路径失败时的最后一级回退）。
|
||||
|
||||
定位与边界：
|
||||
- baostock 是 EOD 数据源，当日数据约 17:30 后才可查，**盘中实时能力为零**；
|
||||
因此本模块只承接 DAY/WEEK/MONTH 的历史 K 线，分钟线/分时/逐笔/板块/实时
|
||||
报价一律返回不可用（None），由上层维持原错误。
|
||||
- 启用条件自动判断：已安装 baostock 且未设置环境变量 ``EASY_TDX_BAOSTOCK=0``
|
||||
即启用；未安装时本模块整体静默关闭，核心功能零影响。
|
||||
- baostock 客户端是单条全局连接且非线程安全，本模块内部全程持锁串行，
|
||||
供 async 调用方经 ``asyncio.to_thread`` 使用。
|
||||
- 数据口径：volume 为股（与 /bars 输出契约一致，无需换算）；停牌日
|
||||
（tradestatus=0 或 volume=0）剔除，与通达信 K 线不含停牌日的口径对齐；
|
||||
复权经 adjustflag 原生支持（QFQ/HFQ/NONE），North Exchange（BJ）不覆盖。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import threading
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pandas as pd
|
||||
|
||||
BAOSTOCK_DISABLE_ENV = "EASY_TDX_BAOSTOCK"
|
||||
|
||||
# baostock 的全局连接锁（该库单连接、非线程安全）
|
||||
_bs_lock = threading.Lock()
|
||||
_logged_in = False
|
||||
# 兜底路径的锁等待上限：拿不到锁说明另一个兜底请求正在进行，
|
||||
# 与其排队不如放弃本次兜底（回退路径宁快勿堵）。
|
||||
_LOCK_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# 支持兜底的周期（baostock frequency）：日线及以上；分钟线/季年线不兜
|
||||
_FREQ_BY_CATEGORY: dict[str, str] = {"DAY": "d", "WEEK": "w", "MONTH": "m"}
|
||||
# 复权映射：baostock adjustflag — 1=后复权 2=前复权 3=不复权
|
||||
_ADJUST_FLAG = {"NONE": "3", "QFQ": "2", "HFQ": "1"}
|
||||
_MARKET_PREFIX = {"SZ": "sz", "SH": "sh"} # BJ baostock 不覆盖
|
||||
|
||||
# 拉取窗口的日历天数系数（start+count 根 × 周期占的日历天 + 节假日缓冲）
|
||||
_WINDOW_DAYS = {"d": (1.6, 30), "w": (7.2, 40), "m": (32.0, 100)}
|
||||
# 偏移窗口规模上限（/bars 的 start 无上界，防极端参数把兜底源拖死）
|
||||
_MAX_TOTAL_BARS = 10_000
|
||||
|
||||
|
||||
def is_enabled() -> bool:
|
||||
"""自动判断兜底是否可用：未禁用（环境变量）且 baostock 已安装。"""
|
||||
disabled = os.environ.get(BAOSTOCK_DISABLE_ENV, "").strip().lower() in {"0", "false", "off"}
|
||||
if disabled:
|
||||
return False
|
||||
try:
|
||||
importlib.import_module("baostock")
|
||||
except ImportError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _login_if_needed(bs: object) -> None:
|
||||
"""确保 baostock 已登录（匿名账户；调用方需已持有 _bs_lock）。"""
|
||||
global _logged_in
|
||||
if _logged_in:
|
||||
return
|
||||
lg = bs.login() # type: ignore[attr-defined]
|
||||
error_code = str(getattr(lg, "error_code", ""))
|
||||
if error_code != "0":
|
||||
raise RuntimeError(f"baostock 登录失败: {getattr(lg, 'error_msg', error_code)}")
|
||||
_logged_in = True
|
||||
|
||||
|
||||
def _query_rows(bs: object, **kwargs: str) -> list[list[str]]:
|
||||
"""执行 query_history_k_data_plus 并取回全部行（调用方需已持锁）。"""
|
||||
rs = bs.query_history_k_data_plus(**kwargs) # type: ignore[attr-defined]
|
||||
if str(getattr(rs, "error_code", "")) != "0":
|
||||
raise RuntimeError(f"baostock 查询失败: {getattr(rs, 'error_msg', '')}")
|
||||
rows: list[list[str]] = []
|
||||
while rs.next() or False:
|
||||
rows.append(rs.get_row_data())
|
||||
return rows
|
||||
|
||||
|
||||
def fetch_bars(
|
||||
market: str,
|
||||
code: str,
|
||||
category: str,
|
||||
start: int,
|
||||
count: int,
|
||||
adjust: str,
|
||||
) -> pd.DataFrame | None:
|
||||
"""拉取日线及以上 K 线，输出对齐 /bars 契约的 DataFrame。
|
||||
|
||||
Args:
|
||||
market: "SZ" / "SH"（BJ 不支持，返回 None）。
|
||||
code: 6 位代码。
|
||||
category: 周期名（DAY/WEEK/MONTH 之外返回 None）。
|
||||
start: 跳过最新 start 根（与 TDX offset 语义一致）。
|
||||
count: 最多返回 count 根。
|
||||
adjust: "NONE" / "QFQ" / "HFQ"。
|
||||
|
||||
Returns:
|
||||
按 [date, open, close, high, low, vol, amount] 列序、时间升序的
|
||||
DataFrame；兜底不可用 / 不适用 / 无数据时返回 None（调用方继续
|
||||
维持原错误，不吞异常）。
|
||||
"""
|
||||
global _logged_in
|
||||
if not is_enabled():
|
||||
return None
|
||||
prefix = _MARKET_PREFIX.get(market.upper())
|
||||
frequency = _FREQ_BY_CATEGORY.get(category.upper())
|
||||
adjustflag = _ADJUST_FLAG.get(adjust.upper())
|
||||
if prefix is None or frequency is None or adjustflag is None:
|
||||
return None
|
||||
|
||||
total = start + count
|
||||
if total <= 0 or total > _MAX_TOTAL_BARS:
|
||||
return None
|
||||
coef, buffer_days = _WINDOW_DAYS[frequency]
|
||||
end_date = datetime.now()
|
||||
start_date = end_date - timedelta(days=total * coef + buffer_days)
|
||||
|
||||
# baostock 全局单连接：持锁串行；等待超时则放弃本次兜底
|
||||
if not _bs_lock.acquire(timeout=_LOCK_TIMEOUT_SECONDS):
|
||||
return None
|
||||
try:
|
||||
bs = importlib.import_module("baostock")
|
||||
try:
|
||||
_login_if_needed(bs)
|
||||
rows = _query_rows(
|
||||
bs,
|
||||
code=f"{prefix}.{code}",
|
||||
fields="date,open,high,low,close,volume,amount,tradestatus",
|
||||
start_date=start_date.strftime("%Y-%m-%d"),
|
||||
end_date=end_date.strftime("%Y-%m-%d"),
|
||||
frequency=frequency,
|
||||
adjustflag=adjustflag,
|
||||
)
|
||||
except Exception:
|
||||
# 连接可能中途断开：重置登录态，下次兜底重新登录
|
||||
_logged_in = False
|
||||
raise
|
||||
except Exception:
|
||||
# 兜底源自身的任何失败都不向上抛：调用方按"无兜底数据"处理
|
||||
return None
|
||||
finally:
|
||||
_bs_lock.release()
|
||||
|
||||
if not rows:
|
||||
return None
|
||||
df = pd.DataFrame(
|
||||
rows, columns=["date", "open", "high", "low", "close", "vol", "amount", "tradestatus"]
|
||||
)
|
||||
for col in ("open", "high", "low", "close", "vol", "amount"):
|
||||
df[col] = pd.to_numeric(df[col], errors="coerce")
|
||||
# 停牌日剔除（tradestatus=0 或无成交），对齐通达信 K 线不含停牌日的口径
|
||||
if "tradestatus" in df.columns:
|
||||
df = df[df["tradestatus"] != "0"]
|
||||
df = df.dropna(subset=["close"])
|
||||
df = df[df["close"] > 0]
|
||||
df = df[df["vol"] > 0]
|
||||
if df.empty:
|
||||
return None
|
||||
df["date"] = pd.to_datetime(df["date"]).dt.normalize()
|
||||
|
||||
# TDX offset 语义：跳过最新 start 根，再取至多 count 根（时间升序）
|
||||
end_pos = len(df) - start
|
||||
if end_pos <= 0:
|
||||
return None
|
||||
df = df.iloc[max(0, end_pos - count) : end_pos]
|
||||
if df.empty:
|
||||
return None
|
||||
|
||||
return df[["date", "open", "close", "high", "low", "vol", "amount"]].reset_index(drop=True)
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
@@ -18,7 +19,7 @@ from easy_tdx.web.convert import (
|
||||
period_times_from_category,
|
||||
)
|
||||
from easy_tdx.web.deps import get_client, get_mac_client_optional
|
||||
from easy_tdx.web.schemas import DataFrameResponse
|
||||
from easy_tdx.web.schemas import BarsResponse, DataFrameResponse
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -216,7 +217,40 @@ async def _fetch_120m(
|
||||
return _resample_pairs(df, count)
|
||||
|
||||
|
||||
@router.get("/bars", response_model=DataFrameResponse)
|
||||
async def _baostock_last_resort(
|
||||
market: str, code: str, category: str, start: int, count: int, adjust: str
|
||||
) -> tuple[pd.DataFrame | None, str | None]:
|
||||
"""TDX 全部路径失败/为空后的最后一级兜底：baostock（仅日线及以上）。
|
||||
|
||||
未安装 baostock / 设置了 EASY_TDX_BAOSTOCK=0 / 周期不适用 / 查询失败
|
||||
一律返回 ``(None, None)``——兜底源自身的任何失败都不影响原错误语义。
|
||||
baostock 客户端阻塞且非线程安全：丢线程池执行，模块内部持锁串行。
|
||||
"""
|
||||
from easy_tdx.sources import baostock as baostock_source
|
||||
|
||||
if not baostock_source.is_enabled():
|
||||
return None, None
|
||||
try:
|
||||
df = await asyncio.to_thread(
|
||||
baostock_source.fetch_bars, market, code, category, start, count, adjust
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — 兜底失败不改变原错误路径
|
||||
_logger.warning("/bars baostock 兜底异常 (%s%s): %s", market, code, exc)
|
||||
return None, None
|
||||
if df is None or df.empty:
|
||||
return None, None
|
||||
_logger.info("/bars 已启用 baostock 兜底 (%s%s %s，%d 根)", market, code, category, len(df))
|
||||
return df, "baostock"
|
||||
|
||||
|
||||
def _bars_resp(df: pd.DataFrame | None, source: str | None) -> BarsResponse:
|
||||
"""构建带来源标注的 K 线响应（source 非 None = 命中兜底源）。"""
|
||||
resp = BarsResponse.from_dataframe(df)
|
||||
resp.source = source
|
||||
return resp
|
||||
|
||||
|
||||
@router.get("/bars", response_model=BarsResponse)
|
||||
async def security_bars(
|
||||
market: str = Query(..., description="市场: SZ, SH, BJ"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
@@ -237,11 +271,14 @@ async def security_bars(
|
||||
),
|
||||
mac_client: Any = Depends(get_mac_client_optional),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
) -> BarsResponse:
|
||||
"""获取股票K线数据（MAC 协议，支持复权）。
|
||||
|
||||
优先走 AsyncMacClient.get_stock_kline（支持 NONE/QFQ/HFQ 复权 + QFQ 负价兜底）；
|
||||
MAC 主机未连接时自动回退 AsyncTdxClient.get_security_bars（无复权，adjust 参数忽略）。
|
||||
多级自动回退：MAC 优先（支持 NONE/QFQ/HFQ 复权 + QFQ 负价兜底）→
|
||||
失败/为空转标准 TdxClient（无复权，adjust 参数忽略）→ 仍失败/为空且
|
||||
周期为日线及以上时，最后一级自动兜底 baostock（需 ``pip install
|
||||
easy-tdx[baostock]``，可用 ``EASY_TDX_BAOSTOCK=0`` 关闭）。兜底命中时
|
||||
响应带 ``source: "baostock"``，否则该字段为 null。
|
||||
输出契约与旧版一致：日线返回 ``date`` 列，分钟线返回 ``datetime`` 列。
|
||||
|
||||
``category=MIN_120`` 为 120 分钟线：MAC 原生 ``Period.MINS × times=120``
|
||||
@@ -256,11 +293,16 @@ async def security_bars(
|
||||
成交量/100，回退路径（标准 TdxClient）已 ×100 还原为股。
|
||||
"""
|
||||
if category.upper() in _MIN_120_ALIASES:
|
||||
df = await _fetch_120m(market, code, start, count, adjust, bar_time, mac_client, client)
|
||||
return _df_resp(_attach_derived(df))
|
||||
df120 = await _fetch_120m(market, code, start, count, adjust, bar_time, mac_client, client)
|
||||
return _bars_resp(_attach_derived(df120), None)
|
||||
|
||||
cat = category_from_str(category)
|
||||
df: pd.DataFrame | None = None
|
||||
source: str | None = None
|
||||
last_exc: Exception | None = None
|
||||
|
||||
if mac_client is not None:
|
||||
try:
|
||||
period, times = period_times_from_category(cat)
|
||||
df = await mac_client.get_stock_kline(
|
||||
market_value_from_str(market),
|
||||
@@ -274,19 +316,42 @@ async def security_bars(
|
||||
)
|
||||
# daily_plus：日线及以上周期 datetime→date（枚举值无序，显式查表判定）
|
||||
df = _normalize_mac_df(df, daily_plus=_is_daily_plus(cat))
|
||||
else:
|
||||
# MAC 不可用：回退标准 TdxClient（无复权），adjust 参数忽略
|
||||
except Exception as exc: # noqa: BLE001 — 降级到标准客户端，不中断
|
||||
last_exc = exc
|
||||
df = None
|
||||
_logger.warning("/bars MAC 获取失败，转标准 TdxClient (%s%s): %s", market, code, exc)
|
||||
if df is None or df.empty:
|
||||
if mac_client is None:
|
||||
_logger.warning(
|
||||
"/bars MAC 客户端未连接，回退标准 TdxClient（不支持复权，adjust=%s 被忽略）",
|
||||
adjust,
|
||||
)
|
||||
elif df is not None and df.empty:
|
||||
# MAC 抛异常的情况已在 except 分支记录
|
||||
_logger.info("/bars MAC 返回空，转标准 TdxClient (%s%s)", market, code)
|
||||
try:
|
||||
df = await client.get_security_bars(
|
||||
market_from_str(market), code, cat, start, count, bar_time=bar_time
|
||||
)
|
||||
return _df_resp(_attach_derived(df))
|
||||
except Exception as exc: # noqa: BLE001 — 降级到 baostock，不中断
|
||||
last_exc = exc
|
||||
df = None
|
||||
_logger.warning("/bars 标准 TdxClient 获取失败 (%s%s): %s", market, code, exc)
|
||||
|
||||
if df is None or df.empty:
|
||||
bdf, bsource = await _baostock_last_resort(market, code, category, start, count, adjust)
|
||||
if bdf is not None:
|
||||
df, source = bdf, bsource
|
||||
|
||||
if df is None:
|
||||
# TDX 两级都抛了异常且兜底不可用：维持原错误语义（503/500）
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
df = pd.DataFrame()
|
||||
return _bars_resp(_attach_derived(df), source)
|
||||
|
||||
|
||||
@router.get("/bars/index", response_model=DataFrameResponse)
|
||||
@router.get("/bars/index", response_model=BarsResponse)
|
||||
async def index_bars(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
@@ -297,19 +362,44 @@ async def index_bars(
|
||||
"start", description="时间戳: start=bar开始时间(默认) / end=bar结束时间(对齐Tushare)"
|
||||
),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
) -> BarsResponse:
|
||||
"""获取指数K线数据。
|
||||
|
||||
指数K线并非所有 TDX 服务器都提供：失败/为空时自动兜底 baostock
|
||||
（仅日线及以上，见 /bars 说明），命中时响应带 ``source: "baostock"``。
|
||||
|
||||
vol 单位：日线/周线/月线/季线/年线 = 成交量(手)（周及以上周期服务端
|
||||
原样返回真实成交量/100，已 ×100 还原）；**分钟线协议不提供成交量**
|
||||
（报文中该字段实为成交额/100），vol 为 ``null``，请勿当作成交量使用。
|
||||
|
||||
每根 bar 同样附带 ``pre_close/change/change_pct/amplitude_pct`` 衍生字段。
|
||||
"""
|
||||
df: pd.DataFrame | None = None
|
||||
source: str | None = None
|
||||
last_exc: Exception | None = None
|
||||
try:
|
||||
df = await client.get_index_bars(
|
||||
market_from_str(market), code, category_from_str(category), start, count, bar_time=bar_time
|
||||
market_from_str(market),
|
||||
code,
|
||||
category_from_str(category),
|
||||
start,
|
||||
count,
|
||||
bar_time=bar_time,
|
||||
)
|
||||
return _df_resp(_attach_derived(df))
|
||||
except Exception as exc: # noqa: BLE001 — 降级到 baostock，不中断
|
||||
last_exc = exc
|
||||
_logger.warning("/bars/index TdxClient 获取失败 (%s%s): %s", market, code, exc)
|
||||
|
||||
if df is None or df.empty:
|
||||
bdf, bsource = await _baostock_last_resort(market, code, category, start, count, "QFQ")
|
||||
if bdf is not None:
|
||||
df, source = bdf, bsource
|
||||
|
||||
if df is None:
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
df = pd.DataFrame()
|
||||
return _bars_resp(_attach_derived(df), source)
|
||||
|
||||
|
||||
@router.get("/minute", response_model=DataFrameResponse)
|
||||
|
||||
@@ -124,6 +124,18 @@ class DataFrameResponse(BaseModel):
|
||||
return cls(data=[], count=0)
|
||||
|
||||
|
||||
class BarsResponse(DataFrameResponse):
|
||||
"""K 线响应。``source`` 非 None 表示数据来自自动兜底源（如 baostock，
|
||||
TDX 全部路径失败时启用）——口径透明：调用方可据此展示数据来源。"""
|
||||
|
||||
source: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dataframe(cls, df: Any) -> BarsResponse:
|
||||
resp = DataFrameResponse.from_dataframe(df)
|
||||
return cls(data=resp.data, count=resp.count)
|
||||
|
||||
|
||||
class DictResponse(BaseModel):
|
||||
"""通用 dict 响应（用于非 DataFrame 返回值）。"""
|
||||
|
||||
|
||||
@@ -0,0 +1,284 @@
|
||||
"""baostock 自动兜底数据源单测（离线，注入假 baostock 模块）。
|
||||
|
||||
覆盖：参数映射（代码/周期/复权）、offset 切片语义、停牌日剔除、
|
||||
可用性门控（环境变量 / 未安装）、/bars 与 /bars/index 的端到端兜底、
|
||||
TDX 正常时绝不触发兜底。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 假 baostock 模块
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeLoginResult:
|
||||
error_code = "0"
|
||||
error_msg = ""
|
||||
|
||||
|
||||
class _FakeResultData:
|
||||
def __init__(self, rows: list[list[str]]):
|
||||
self._rows = rows
|
||||
self._i = 0
|
||||
self.error_code = "0"
|
||||
self.error_msg = ""
|
||||
|
||||
def next(self) -> bool:
|
||||
if self._i < len(self._rows):
|
||||
self._i += 1
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_row_data(self) -> list[str]:
|
||||
return self._rows[self._i - 1]
|
||||
|
||||
|
||||
def _fake_rows(n: int, end: str = "2026-09-04") -> list[list[str]]:
|
||||
"""n 个交易日的日线行：date, open, high, low, close, volume, amount, tradestatus。"""
|
||||
dates = pd.bdate_range(end=end, periods=n).strftime("%Y-%m-%d")
|
||||
return [[d, "10.0", "11.0", "9.5", "10.5", "100000", "1050000.0", "1"] for d in dates]
|
||||
|
||||
|
||||
def _install_fake_bs(
|
||||
rows: list[list[str]] | None,
|
||||
captured: dict,
|
||||
*,
|
||||
query_error: bool = False,
|
||||
) -> types.ModuleType:
|
||||
mod = types.ModuleType("baostock")
|
||||
|
||||
def _login(): # type: ignore[no-untyped-def]
|
||||
captured["login"] = captured.get("login", 0) + 1
|
||||
return _FakeLoginResult()
|
||||
|
||||
mod.login = _login # type: ignore[attr-defined]
|
||||
mod.logout = lambda: None # type: ignore[attr-defined]
|
||||
|
||||
def query_history_k_data_plus(**kwargs): # type: ignore[no-untyped-def]
|
||||
captured.update(kwargs)
|
||||
captured["calls"] = captured.get("calls", 0) + 1
|
||||
if query_error:
|
||||
result = _FakeResultData([])
|
||||
result.error_code = "10001"
|
||||
result.error_msg = "网络异常"
|
||||
return result
|
||||
return _FakeResultData(rows or [])
|
||||
|
||||
mod.query_history_k_data_plus = query_history_k_data_plus # type: ignore[attr-defined]
|
||||
sys.modules["baostock"] = mod
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fake_bs(monkeypatch: pytest.MonkeyPatch):
|
||||
"""注入假模块 + 复位模块级登录态；测试结束移除。"""
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(bs_source, "_logged_in", False)
|
||||
monkeypatch.delenv(bs_source.BAOSTOCK_DISABLE_ENV, raising=False)
|
||||
_install_fake_bs(_fake_rows(10), captured)
|
||||
yield captured
|
||||
sys.modules.pop("baostock", None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 源模块行为
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fetch_maps_args_and_matches_contract(fake_bs):
|
||||
"""代码/周期/复权映射正确；输出列序与 vol 单位（股，不换算）符合 /bars 契约。"""
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
df = bs_source.fetch_bars("SH", "600519", "DAY", 0, 5, "QFQ")
|
||||
assert df is not None and len(df) == 5
|
||||
assert list(df.columns) == ["date", "open", "close", "high", "low", "vol", "amount"]
|
||||
assert fake_bs["code"] == "sh.600519"
|
||||
assert fake_bs["frequency"] == "d"
|
||||
assert fake_bs["adjustflag"] == "2" # QFQ
|
||||
# 时间升序，最后一根是最新交易日
|
||||
assert df["date"].iloc[-1] == pd.Timestamp("2026-09-04")
|
||||
assert (df["vol"] == 100000).all() # baostock volume=股，与 /bars 契约一致，不换算
|
||||
|
||||
|
||||
def test_offset_slice_matches_tdx_semantics(fake_bs):
|
||||
"""start=跳过最新 N 根：30 根里 start=5, count=10 → 返回第 16~25 根。"""
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
_install_fake_bs(_fake_rows(30), fake_bs)
|
||||
df = bs_source.fetch_bars("SZ", "000001", "DAY", 5, 10, "QFQ")
|
||||
assert df is not None and len(df) == 10
|
||||
dates = df["date"].dt.strftime("%Y-%m-%d").tolist()
|
||||
expected = pd.bdate_range(end="2026-09-04", periods=30).strftime("%Y-%m-%d").tolist()
|
||||
assert dates[0] == expected[15]
|
||||
assert dates[-1] == expected[24]
|
||||
|
||||
|
||||
def test_suspension_rows_dropped(fake_bs):
|
||||
"""停牌日（tradestatus=0 / volume=0）剔除，对齐通达信 K 线口径。"""
|
||||
rows = _fake_rows(6)
|
||||
rows[2] = [rows[2][0], "0", "0", "0", "0", "0", "0", "0"] # 停牌日
|
||||
_install_fake_bs(rows, fake_bs)
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
df = bs_source.fetch_bars("SZ", "000001", "DAY", 0, 10, "QFQ")
|
||||
assert df is not None and len(df) == 5
|
||||
assert (df["vol"] > 0).all()
|
||||
|
||||
|
||||
def test_disabled_via_env(fake_bs, monkeypatch: pytest.MonkeyPatch):
|
||||
"""EASY_TDX_BAOSTOCK=0 显式关闭：不安装也不调用。"""
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
monkeypatch.setenv(bs_source.BAOSTOCK_DISABLE_ENV, "0")
|
||||
assert bs_source.is_enabled() is False
|
||||
assert bs_source.fetch_bars("SH", "600519", "DAY", 0, 5, "QFQ") is None
|
||||
assert "login" not in fake_bs
|
||||
|
||||
|
||||
def test_missing_module_returns_none(monkeypatch: pytest.MonkeyPatch):
|
||||
"""未安装 baostock：静默返回 None（兜底环自动关闭）。"""
|
||||
monkeypatch.delenv("EASY_TDX_BAOSTOCK", raising=False)
|
||||
monkeypatch.setitem(sys.modules, "baostock", None) # import 时抛 ImportError
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
assert bs_source.is_enabled() is False
|
||||
assert bs_source.fetch_bars("SH", "600519", "DAY", 0, 5, "QFQ") is None
|
||||
|
||||
|
||||
def test_unsupported_inputs(fake_bs):
|
||||
"""BJ 市场 / 分钟线周期 / 非法复权 / 超大窗口：不适用即 None。"""
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
assert bs_source.fetch_bars("BJ", "430047", "DAY", 0, 5, "QFQ") is None
|
||||
assert bs_source.fetch_bars("SH", "600519", "MIN_5", 0, 5, "QFQ") is None
|
||||
assert bs_source.fetch_bars("SH", "600519", "SEASON", 0, 5, "QFQ") is None
|
||||
assert bs_source.fetch_bars("SH", "600519", "DAY", 0, 5, "FOO") is None
|
||||
assert bs_source.fetch_bars("SH", "600519", "DAY", 99999, 800, "QFQ") is None
|
||||
assert "calls" not in fake_bs
|
||||
|
||||
|
||||
def test_query_error_returns_none(fake_bs):
|
||||
"""baostock 查询失败：返回 None 且不向上抛（兜底失败不改变原错误路径）。"""
|
||||
_install_fake_bs([], fake_bs, query_error=True)
|
||||
from easy_tdx.sources import baostock as bs_source
|
||||
|
||||
assert bs_source.fetch_bars("SH", "600519", "DAY", 0, 5, "QFQ") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /bars 与 /bars/index 端到端兜底
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _bars_app(mac_client, tdx_client):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from easy_tdx.web.errors import register_exception_handlers
|
||||
from easy_tdx.web.routers import bars
|
||||
|
||||
app = FastAPI()
|
||||
register_exception_handlers(app)
|
||||
app.include_router(bars.router, prefix="/api/v1")
|
||||
app.state.tdx_client = tdx_client
|
||||
app.state.mac_client = mac_client
|
||||
return app
|
||||
|
||||
|
||||
class _RaisingMac:
|
||||
async def get_stock_kline(self, *args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise RuntimeError("MAC 连接失败")
|
||||
|
||||
|
||||
class _RaisingTdx:
|
||||
async def get_security_bars(self, *args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise RuntimeError("标准协议连接失败")
|
||||
|
||||
async def get_index_bars(self, *args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise RuntimeError("标准协议连接失败")
|
||||
|
||||
|
||||
class _OkMac:
|
||||
async def get_stock_kline(self, *args, **kwargs): # noqa: ANN002, ANN003
|
||||
return pd.DataFrame(
|
||||
{
|
||||
"datetime": pd.bdate_range(end="2026-09-04", periods=5),
|
||||
"open": [10.0] * 5,
|
||||
"close": [10.5] * 5,
|
||||
"high": [11.0] * 5,
|
||||
"low": [9.5] * 5,
|
||||
"vol": [100000] * 5,
|
||||
"amount": [1050000.0] * 5,
|
||||
"float_shares": [0.0] * 5,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_bars_endpoint_falls_back_to_baostock(fake_bs, monkeypatch: pytest.MonkeyPatch):
|
||||
"""MAC 与标准协议都失败 → baostock 兜底命中，响应带 source 字段。"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
_install_fake_bs(_fake_rows(10), fake_bs)
|
||||
with TestClient(_bars_app(_RaisingMac(), _RaisingTdx())) as client:
|
||||
resp = client.get("/api/v1/bars", params={"market": "SH", "code": "600519"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["source"] == "baostock"
|
||||
assert body["count"] == 10
|
||||
assert "date" in body["data"][0]
|
||||
assert "change_pct" in body["data"][0]
|
||||
|
||||
|
||||
def test_bars_endpoint_tdx_ok_never_calls_baostock(fake_bs):
|
||||
"""TDX 正常出数时兜底绝不触发：source 为 None，baostock 零调用。"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
with TestClient(_bars_app(_OkMac(), _RaisingTdx())) as client:
|
||||
resp = client.get("/api/v1/bars", params={"market": "SH", "code": "600519"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["source"] is None
|
||||
assert body["count"] == 5
|
||||
assert "login" not in fake_bs
|
||||
|
||||
|
||||
def test_bars_endpoint_no_fallback_available_keeps_error(fake_bs, monkeypatch: pytest.MonkeyPatch):
|
||||
"""TDX 全败且兜底不可用：维持原错误语义（500），不返回空数据伪装成功。"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
monkeypatch.delenv("EASY_TDX_BAOSTOCK", raising=False)
|
||||
monkeypatch.setitem(sys.modules, "baostock", None)
|
||||
# raise_server_exceptions=False：模拟生产环境由服务端中间件返回 500
|
||||
with TestClient(
|
||||
_bars_app(_RaisingMac(), _RaisingTdx()), raise_server_exceptions=False
|
||||
) as client:
|
||||
resp = client.get("/api/v1/bars", params={"market": "SH", "code": "600519"})
|
||||
assert resp.status_code == 500
|
||||
assert "连接失败" in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_index_endpoint_falls_back_to_baostock(fake_bs):
|
||||
"""/bars/index：TDX 失败 → baostock 兜底（指数代码同格式）。"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
_install_fake_bs(_fake_rows(10), fake_bs)
|
||||
with TestClient(_bars_app(None, _RaisingTdx())) as client:
|
||||
resp = client.get(
|
||||
"/api/v1/bars/index", params={"market": "SH", "code": "000001", "category": "DAY"}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["source"] == "baostock"
|
||||
assert fake_bs["code"] == "sh.000001"
|
||||
Reference in New Issue
Block a user