feat(bs): --backfill-valuation估值历史窗回补入口——只拉9列不动dbbardata,断点buffer先落盘done后写,终局merge原子写回 [vps]
This commit is contained in:
@@ -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")
|
||||
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user