From 150465a90eb3654af3e39071e9936621e690ce08 Mon Sep 17 00:00:00 2001 From: claude_dev Date: Tue, 8 Sep 2026 20:57:03 +0800 Subject: [PATCH] =?UTF-8?q?feat(bs):=20--backfill-valuation=E4=BC=B0?= =?UTF-8?q?=E5=80=BC=E5=8E=86=E5=8F=B2=E7=AA=97=E5=9B=9E=E8=A1=A5=E5=85=A5?= =?UTF-8?q?=E5=8F=A3=E2=80=94=E2=80=94=E5=8F=AA=E6=8B=899=E5=88=97?= =?UTF-8?q?=E4=B8=8D=E5=8A=A8dbbardata,=E6=96=AD=E7=82=B9buffer=E5=85=88?= =?UTF-8?q?=E8=90=BD=E7=9B=98done=E5=90=8E=E5=86=99,=E7=BB=88=E5=B1=80merg?= =?UTF-8?q?e=E5=8E=9F=E5=AD=90=E5=86=99=E5=9B=9E=20[vps]?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/data_platform/bs_eod.py | 245 +++++++++++++++++- .../test_bs_eod_valuation_backfill.py | 150 +++++++++++ 2 files changed, 389 insertions(+), 6 deletions(-) create mode 100644 tests/data_platform/test_bs_eod_valuation_backfill.py diff --git a/scripts/data_platform/bs_eod.py b/scripts/data_platform/bs_eod.py index 2af06bb..56e0d59 100644 --- a/scripts/data_platform/bs_eod.py +++ b/scripts/data_platform/bs_eod.py @@ -217,6 +217,26 @@ def fetch_all_stocks_with_timeout(timeout=120): return _with_timeout(fetch_all_stocks, timeout=timeout) +def _valuation_frame(code, prefix, rows, fields): + """估值 K 线行 -> 10 列 vdf [symbol,exchange,date,peTTM,psTTM,pcfNcfTTM, + pbMRQ,turn,pctChg,isST]。 + + fields 为请求字段串(列名对齐即可, 顺序无关): upsert_daily 传 DAILY_FIELDS + (15 列含 OHLCV), 回补传 VAL_BACKFILL_FIELDS (9 列纯估值)。 + 空串 peTTM/pbMRQ -> NaN (亏损/净资负), isST 字符串 -> int。 + """ + df = pd.DataFrame(rows, columns=fields.split(",")) + for c in ["turn", "pctChg", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ"]: + if c in df.columns: + df[c] = pd.to_numeric(df[c], errors="coerce") + vdf = df[["date", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ", "turn", + "pctChg", "isST"]].copy() + vdf["isST"] = pd.to_numeric(vdf["isST"], errors="coerce").fillna(0).astype(int) + vdf.insert(0, "symbol", code) + vdf.insert(1, "exchange", EXC_MAP[prefix]) + return vdf + + def upsert_daily(conn, code, prefix, rows): """日线 rows -> dbbardata('d') + valuation_baostock 当年 parquet 追加。""" if not rows: @@ -241,10 +261,7 @@ def upsert_daily(conn, code, prefix, rows): "open_price,high_price,low_price,close_price) VALUES (?,?,?,?,?,?,?,?,?,?,?)", db.itertuples(index=False, name=None)) # pe/pb -> parquet 追加 (isST->int, 修 pyarrow ArrowTypeError) - vdf = df[["date", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ", "turn", "pctChg", "isST"]].copy() - vdf["isST"] = pd.to_numeric(vdf["isST"], errors="coerce").fillna(0).astype(int) - vdf.insert(0, "symbol", code) - vdf.insert(1, "exchange", exc) + vdf = _valuation_frame(code, prefix, rows, DAILY_FIELDS) yr = dt.date.today().year p = VAL_DIR / f"{yr}.parquet" if p.exists(): @@ -371,14 +388,230 @@ def _process_one_stock(conn, code, prefix, args, start, end): return n1, n2 -def main(): +# ======================== valuation_baostock 历史窗回补 (spec §19.11-a) ======================== +# 2026.parquet 仅 08-13 起(bs 日喂起点), H1 缺口靠基线层 em valuation 异源口径补。 +# 回补=按 (股, 窗口) 只拉估值 9 列 -> 中间产物落 data/_backfill_valuation/{tag}/ +# (独立目录, 不进 sync_valuation_daily 的 scp -r 整目录拉取面) -> 终局与年文件 +# merge 原子写回。断点序=buffer 先落盘、done 后写: 崩在中间最多整批重拉(键唯一幂等), 行永不丢。 + +VAL_BACKFILL_FIELDS = ("date,code,turn,pctChg,peTTM,psTTM," + "pcfNcfTTM,pbMRQ,isST") # 无 OHLCV: 不动 dbbardata +BACKFILL_BATCH = 500 # 只/批: buffer 落盘+done 写入粒度 + + +def merge_valuation_frames(old, new): + """年文件 ∪ 回补段: drop_duplicates(symbol,date,keep='last')(与 upsert_daily + 同语义) + sort。old=None(年文件不存在)时直接排序返回 new 副本。""" + if old is None: + return new.sort_values(["symbol", "date"]).reset_index(drop=True) + out = (pd.concat([old, new], ignore_index=True) + .drop_duplicates(["symbol", "date"], keep="last") + .sort_values(["symbol", "date"]) + .reset_index(drop=True)) + return out + + +def build_backfill_stock_list(current, snapshot): + """当前全列表(全A含退市) ∪ 时点快照(query_all_stock) 并集去重保序。 + + 时点兜底防 query_stock_basic 退市覆盖缺口; snapshot 空/None 只用当前列表。 + """ + seen = set() + out = [] + for item in list(current) + list(snapshot or []): + if item not in seen: + seen.add(item) + out.append(item) + return out + + +def midpoint_date(start, end): + """窗口中点日(str YYYY-MM-DD) — query_all_stock 时点采样日。""" + s = dt.datetime.strptime(start, "%Y-%m-%d").date() + e = dt.datetime.strptime(end, "%Y-%m-%d").date() + return (s + (e - s) // 2).strftime("%Y-%m-%d") + + +def atomic_write_parquet(df, path): + """tmp 写入 + os.replace 原子替换(2-3h 回补成果不能毁于写一半)。""" + tmp = Path(str(path) + ".tmp") + df.to_parquet(tmp, index=False) + os.replace(tmp, path) + + +def _fetch_all_stock_at(day): + """query_all_stock(day) -> [(code, 'sh'/'sz')]; 非交易日/无数据返 []。""" global QUERY_COUNT + QUERY_COUNT += 1 + rs = bs.query_all_stock(day=day) + if rs.error_code != "0": + raise RuntimeError(f"query_all_stock: {rs.error_code} {rs.error_msg}") + idx = {n: i for i, n in enumerate(list(rs.fields))} + out = [] + while rs.next(): + r = rs.get_row_data() + bc = r[idx["code"]] + if "." not in bc: + continue + prefix, num = bc.split(".", 1) + if prefix in ("sh", "sz") and len(num) == 6 and num.isdigit(): + out.append((num, prefix)) + return out + + +def _flush_buffer(buffer_path, frames): + """批内 frames 并入 buffer.parquet(键唯一, 重拉幂等)。""" + new = (pd.concat(frames, ignore_index=True) if len(frames) > 1 + else frames[0].reset_index(drop=True)) + if buffer_path.exists(): + old = pd.read_parquet(buffer_path) + new = pd.concat([old, new], ignore_index=True) + new.drop_duplicates(["symbol", "date"], keep="last").sort_values( + ["symbol", "date"]).to_parquet(buffer_path, index=False) + + +def run_valuation_backfill(start, end): + """一次性回补主流程: 拉取(断点续传) -> 终局 merge 原子写回。单年窗口。""" + global QUERY_COUNT + s = dt.datetime.strptime(start, "%Y-%m-%d").date() + e = dt.datetime.strptime(end, "%Y-%m-%d").date() + if s.year != e.year: + raise SystemExit(f"回补窗口须落在单一年份内: {start}~{end}") + tag = f"{s.year}{s.month:02d}{s.day:02d}_{e.year}{e.month:02d}{e.day:02d}" + work = VAL_DIR.parent / "_backfill_valuation" / tag + work.mkdir(parents=True, exist_ok=True) + done_path = work / "done.txt" + buffer_path = work / "buffer.parquet" + target = VAL_DIR / f"{e.year}.parquet" + log.info("[BACKFILL-VAL] window=%s~%s target=%s work=%s", start, end, + target, work) + + if not login_with_retry(): + log.error("[BACKFILL-VAL] 登录失败, exit 2") + sys.exit(2) + try: + try: + stocks = fetch_all_stocks_with_timeout() + except Exception as e_: + log.error("[BACKFILL-VAL] fetch_all: %s", e_) + sys.exit(1) + snap_day = midpoint_date(start, end) + try: + snapshot = _with_timeout(_fetch_all_stock_at, args=(snap_day,), + timeout=120) + except Exception as e_: + log.warning("query_all_stock(%s) 失败(跳过时点兜底): %s", snap_day, e_) + snapshot = [] + stocks = build_backfill_stock_list(stocks, snapshot) + log.info("[BACKFILL-VAL] 名单=全A∪时点(%s): %d 只", snap_day, len(stocks)) + + done = set() + if done_path.exists(): + done = {ln.strip() for ln in done_path.read_text( + encoding="utf-8").splitlines() if ln.strip()} + log.info("[BACKFILL-VAL] 断点续传: 已完成 %d 只", len(done)) + pending = [(c, p) for c, p in stocks if c not in done] + + stats = {"ok": 0, "empty": 0, "failed": 0, "rows": 0} + t0 = time.time() + batch_vdfs = [] # 批内行缓冲 + batch_codes = [] # 批内 code(待 buffer 落盘后写 done) + stopped = False + with open(done_path, "a", encoding="utf-8") as done_f: + for i, (code, prefix) in enumerate(pending): + if QUERY_COUNT >= DAILY_LIMIT: + log.warning("query %d 达防线 %d, 剩余转下次续跑", + QUERY_COUNT, DAILY_LIMIT) + stopped = True + break + bs_code = f"{prefix}.{code}" + try: + rows = fetch_k_with_timeout(bs_code, VAL_BACKFILL_FIELDS, + "d", start, end) + except Exception as e_: + stats["failed"] += 1 + if stats["failed"] <= 5 or stats["failed"] % 100 == 0: + log.warning("%s fetch err: %s", code, e_) + if not relogin(): + log.error("%s relogin 失败, 跳过", code) + time.sleep(BS_INTERVAL) + continue + if rows: + vdf = _valuation_frame(code, prefix, rows, + VAL_BACKFILL_FIELDS) + # 防御: 窗口越界行过滤(接口契约上不会, 双保险) + vdf = vdf[(vdf["date"] >= start) & (vdf["date"] <= end)] + if len(vdf): + batch_vdfs.append(vdf) + stats["rows"] += len(vdf) + stats["ok"] += 1 + else: + stats["empty"] += 1 + batch_codes.append(code) + if (i + 1) % 100 == 0: + log.info("[BACKFILL-VAL] 进度 %d/%d ok=%d empty=%d " + "failed=%d rows=%d q=%d (%.0fs)", + i + 1, len(pending), stats["ok"], stats["empty"], + stats["failed"], stats["rows"], QUERY_COUNT, + time.time() - t0) + if (i + 1) % RELOGIN_EVERY == 0 and not relogin(): + log.warning("周期 relogin 失败, 继续跑(下次被动 relogin 兜底)") + if len(batch_codes) >= BACKFILL_BATCH: + _flush_buffer(buffer_path, batch_vdfs) + done_f.write("\n".join(batch_codes) + "\n") + done_f.flush() + batch_vdfs, batch_codes = [], [] + if i < len(pending) - 1: + time.sleep(BS_INTERVAL) + # 尾批: buffer 先落盘、done 后写(崩在中间=重拉幂等, 行永不丢) + if batch_codes: + if batch_vdfs: + _flush_buffer(buffer_path, batch_vdfs) + done_f.write("\n".join(batch_codes) + "\n") + log.info("[BACKFILL-VAL] 拉取完成 ok=%d empty=%d failed=%d rows=%d " + "query=%d 耗时%.0fs", stats["ok"], stats["empty"], + stats["failed"], stats["rows"], QUERY_COUNT, time.time() - t0) + + # 终局 merge: buffer ∪ 年文件 -> 原子写回 + if buffer_path.exists(): + buf = pd.read_parquet(buffer_path) + old = pd.read_parquet(target) if target.exists() else None + out = merge_valuation_frames(old, buf) + atomic_write_parquet(out, target) + log.info("[BACKFILL-VAL] merge 完成 %s rows=%d (原 %s) 键唯一=%s", + target, len(out), len(old) if old is not None else 0, + out.duplicated(["symbol", "date"]).sum() == 0) + else: + log.warning("[BACKFILL-VAL] buffer 无数据(全空窗/全失败?), 不动 %s", target) + if stopped: + sys.exit(3) + finally: + try: + bs.logout() + except Exception: + pass + + +def _parse_args(argv=None): ap = argparse.ArgumentParser() ap.add_argument("--limit", type=int, default=0) ap.add_argument("--no-15m", action="store_true") ap.add_argument("--no-daily", action="store_true", help="跳日线, 只跑 15min(用于 15min 重灌快)") - args = ap.parse_args() + ap.add_argument("--backfill-valuation", nargs=2, + metavar=("START", "END"), default=None, + help="一次性: valuation_baostock 历史窗回补 " + "YYYY-MM-DD YYYY-MM-DD(只拉估值列, 不动 dbbardata)") + return ap.parse_args(argv) + + +def main(): + global QUERY_COUNT + args = _parse_args() + if args.backfill_valuation: + start, end = args.backfill_valuation + run_valuation_backfill(start, end) + return today = dt.date.today() end = today.strftime("%Y-%m-%d") diff --git a/tests/data_platform/test_bs_eod_valuation_backfill.py b/tests/data_platform/test_bs_eod_valuation_backfill.py new file mode 100644 index 0000000..36e8951 --- /dev/null +++ b/tests/data_platform/test_bs_eod_valuation_backfill.py @@ -0,0 +1,150 @@ +# -*- coding: utf-8 -*- +"""TDD for bs_eod.py valuation_baostock 一次性回补 (--backfill-valuation, spec §19.11-a). + +2026.parquet 仅 08-13 起(bs 日喂起点), H1 缺口靠基线层 em valuation 异源口径补。 +回补=baostock 按 (股, 窗口) 只拉估值 10 列 -> 与现有年文件 merge -> 原子写回。 + +测试覆盖(纯函数, 不依赖真实 baostock): + 1. _valuation_frame: 估值行 -> 10 列 vdf (backfill 字段子集 / DAILY_FIELDS 全列两兼容) + 2. merge_valuation_frames: 零重叠行数=和 / 重叠 new 胜 / 空 old / 键唯一+排序 + 3. build_backfill_stock_list: 当前全列表 ∪ 时点快照 并集去重保序 + 4. midpoint_date: 窗口中点日 + 5. atomic_write_parquet: 内容对 + 无 .tmp 残留 + 6. CLI: --backfill-valuation START END 解析存在 +""" +import sys +import types +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +# Mock baostock before import (Mac 可能没装 / 不依赖网络, 同 test_bs_eod_resilience) +if "baostock" not in sys.modules: + sys.modules["baostock"] = MagicMock() + +from scripts.data_platform import bs_eod # noqa: E402 + + +def _vdf(rows): + """便捷构造: [(symbol, date, peTTM, pbMRQ)] -> 最小 vdf。""" + return pd.DataFrame(rows, columns=["symbol", "date", "peTTM", "pbMRQ"]) + + +# ======================== 1. _valuation_frame ======================== + +class TestValuationFrame: + BACKFILL_FIELDS = ("date,code,turn,pctChg,peTTM,psTTM," + "pcfNcfTTM,pbMRQ,isST") + + def test_backfill_fields_subset(self): + """backfill 只拉估值列(无 OHLCV)也能构造同款 10 列 vdf。""" + rows = [["2026-03-02", "sh.600519", "0.8", "1.5", "25.1", + "10.2", "15.0", "8.1", "0"], + ["2026-03-03", "sh.600519", "0.9", "-0.2", "", "10.3", + "15.1", "8.2", "1"]] + vdf = bs_eod._valuation_frame("600519", "sh", rows, + self.BACKFILL_FIELDS) + assert list(vdf.columns) == ["symbol", "exchange", "date", "peTTM", + "psTTM", "pcfNcfTTM", "pbMRQ", "turn", + "pctChg", "isST"] + assert len(vdf) == 2 + assert (vdf["symbol"] == "600519").all() + assert (vdf["exchange"] == "SSE").all() + # 空串 peTTM -> NaN (亏损/净资负), isST 字符串 -> int + assert pd.isna(vdf.iloc[1]["peTTM"]) + assert vdf["isST"].tolist() == [0, 1] + + def test_daily_fields_full_still_works(self): + """DAILY_FIELDS(15 列含 OHLCV)输入兼容 — upsert_daily 复用不回归。""" + row = ["2026-03-02", "sh.600519", "1700", "1710", "1690", "1705", + "10000", "1705000", "0.8", "1.5", "25.1", "10.2", "15.0", + "8.1", "0"] + vdf = bs_eod._valuation_frame("600519", "sh", [row], + bs_eod.DAILY_FIELDS) + assert vdf.iloc[0]["peTTM"] == 25.1 and vdf.iloc[0]["isST"] == 0 + + +# ======================== 2. merge_valuation_frames ======================== + +class TestMergeValuationFrames: + def test_disjoint_union(self): + """零重叠(H1 窗止 08-12 vs 现有 08-13 起): 行数=和, 键唯一。""" + old = _vdf([("600519", "2026-08-13", 20.0, 8.0), + ("600519", "2026-08-14", 20.1, 8.0)]) + new = _vdf([("600519", "2026-01-05", 22.0, 7.5)]) + out = bs_eod.merge_valuation_frames(old, new) + assert len(out) == 3 + assert out.duplicated(["symbol", "date"]).sum() == 0 + # 按 (symbol, date) 排序 + assert out["date"].tolist() == ["2026-01-05", "2026-08-13", + "2026-08-14"] + + def test_overlap_new_wins(self): + """重叠时 new 胜(与 upsert_daily keep='last' 同语义)。""" + old = _vdf([("600519", "2026-08-13", 20.0, 8.0)]) + new = _vdf([("600519", "2026-08-13", 21.0, 8.1)]) + out = bs_eod.merge_valuation_frames(old, new) + assert len(out) == 1 and out.iloc[0]["peTTM"] == 21.0 + + def test_empty_old_passthrough(self): + """年文件不存在时 old=None: 返回 new 的排序副本。""" + new = _vdf([("000001", "2026-02-02", 9.0, 1.1), + ("000001", "2026-02-01", 8.9, 1.1)]) + out = bs_eod.merge_valuation_frames(None, new) + assert len(out) == 2 + assert out["date"].tolist() == ["2026-02-01", "2026-02-02"] + + def test_no_input_mutation(self): + """不可变: 输入 df 不被就地修改。""" + old = _vdf([("600519", "2026-08-13", 20.0, 8.0)]) + new = _vdf([("600519", "2026-01-05", 22.0, 7.5)]) + old_before, new_before = old.copy(), new.copy() + bs_eod.merge_valuation_frames(old, new) + pd.testing.assert_frame_equal(old, old_before) + pd.testing.assert_frame_equal(new, new_before) + + +# ======================== 3. build_backfill_stock_list ======================== + +class TestBuildBackfillStockList: + def test_union_dedup_order(self): + """并集去重保序, 当前列表优先。""" + current = [("600519", "sh"), ("000001", "sz")] + snapshot = [("000001", "sz"), ("600999", "sh")] # 时点兜底多出的 + out = bs_eod.build_backfill_stock_list(current, snapshot) + assert out == [("600519", "sh"), ("000001", "sz"), ("600999", "sh")] + + def test_snapshot_none(self): + """query_all_stock 时点日非交易日返空 -> 只用当前列表。""" + current = [("600519", "sh")] + assert bs_eod.build_backfill_stock_list(current, None) == current + assert bs_eod.build_backfill_stock_list(current, []) == current + + +# ======================== 4. midpoint_date ======================== + +def test_midpoint_date(): + assert bs_eod.midpoint_date("2026-01-01", "2026-08-12") == "2026-04-22" + assert bs_eod.midpoint_date("2026-06-01", "2026-06-30") == "2026-06-15" + + +# ======================== 5. atomic_write_parquet ======================== + +def test_atomic_write_parquet(tmp_path): + p = tmp_path / "2026.parquet" + df = _vdf([("600519", "2026-01-05", 22.0, 7.5)]) + bs_eod.atomic_write_parquet(df, p) + back = pd.read_parquet(p) + pd.testing.assert_frame_equal(back, df) + assert not (tmp_path / "2026.parquet.tmp").exists() # 无残留 + + +# ======================== 6. CLI 解析 ======================== + +def test_cli_backfill_flag(monkeypatch): + """--backfill-valuation START END 存在且解析为两元组。""" + monkeypatch.setattr(sys, "argv", ["bs_eod.py", "--backfill-valuation", + "2026-01-01", "2026-08-12"]) + args = bs_eod._parse_args() + assert args.backfill_valuation == ["2026-01-01", "2026-08-12"]