sanguo/sanguo_vnpy_v2
2
0
Fork 0
Code Issues 6 Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix(data): 中证1000/2000 历史成份股补全(parse+migrate 聚合)

Browse Source 
parse_csindex_announce.py 3 处修复:
- 932000 launch xlsx header 定位"证券代码"列(原 row[0]=指数代码 bug, distinct=1)
- 000852 列表搜索用 indexCode payload 精准拉(28 -> 96 公告, 回溯到 2016)
- akshare current snapshot header 定位"成分券代码"列(原 iloc[0]=日期 bug, current=0)

migrate_constituent.py:
- 加 _read_announce_union_aggregated(): 000852/932000 用 announce_union 全集
  替换原 snapshot-only 路径, was_removed 治幸存者偏差
- 缺 announce_union 时回退旧 snapshot 逻辑(向后兼容)

TDD: tests/portfolio/test_migrate_announce_union.py 7 测全过(5 unit + 2 integration)
Mac 产出验证:
- 000852: distinct 1220 -> 1672, was_removed 0 -> 672
- 932000: distinct 1 -> 2684(launch 修复), was_removed 0 -> 684
This commit is contained in:
Claude Dev
2026-07-23 23:45:55 +08:00
parent 114a69e997
commit b237c2d7e7
3 changed files with 448 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/data_platform/migrate_constituent.py
+82 -16
View File
@@ -77,6 +77,61 @@ def _read_baostock_constituent(c: sqlite3.Connection) -> pd.DataFrame:
return df
def _read_announce_union_aggregated(hist_dir: str,
index_codes: list[str]) -> pd.DataFrame:
"""读 *_announce_union.parquet, 聚合成全集型(index_code/code/code_name/
in_current/was_removed/source)。
announce_union schema:
updateDate / index_code / code / code_name / adjust_type / notice_id / source
adjust_type: add | remove | current | initial
全集聚合(spec §14):
ever_codes = announce_union 所有 distinct code(任意 adjust_type)
current_codes = announce_union 中 adjust_type='current' 的 code
(∪ <idx>_snapshot.parquet 兜底, 兼容旧 announce 无 current 行)
in_current = code in current_codes
was_removed = not in_current(曾经入选已踢)
code_name = 优先 snapshot 当前名(announce 的历史名可能过时)
缺 announce_union 文件 -> 返回空 DataFrame(上层走 snapshot 兜底)。
"""
cols = ["index_code", "code", "code_name", "in_current", "was_removed", "source"]
rows = []
for idx in index_codes:
ann_path = os.path.join(hist_dir, f"{idx}_announce_union.parquet")
if not os.path.exists(ann_path):
continue
ann = pd.read_parquet(ann_path)
ann["code"] = ann["code"].apply(norm_code)
# current 兜底: announce 内 current 行 + snapshot 文件
current_codes = set(ann.loc[ann["adjust_type"] == "current", "code"])
snap_path = os.path.join(hist_dir, f"{idx}_snapshot.parquet")
snap_name_map: dict[str, str] = {}
if os.path.exists(snap_path):
snap = pd.read_parquet(snap_path)
snap["code"] = snap["code"].apply(norm_code)
current_codes |= set(snap["code"])
snap_name_map = dict(zip(snap["code"], snap["code_name"]))
# ever = distinct(code), name 取 announce 或 snapshot(优先 snapshot)
ever = (
ann[["code", "code_name"]]
.drop_duplicates(subset=["code"])
)
for _, r in ever.iterrows():
code = r["code"]
in_cur = code in current_codes
name = snap_name_map.get(code) or r["code_name"]
rows.append({
"index_code": idx, "code": code, "code_name": name,
"in_current": int(in_cur), "was_removed": int(not in_cur),
"source": "csindex_announce",
})
if not rows:
return pd.DataFrame(columns=cols)
return pd.DataFrame(rows, columns=cols)
def migrate(db_path: str) -> None:
"""(幂等)重建 constituent_unified_staging。
@@ -115,23 +170,30 @@ def migrate(db_path: str) -> None:
df_deep["code"] = df_deep["code"].apply(norm_code)
print(f"[深证 union] rows={len(df_deep)}")
# 3. 中证 snapshot
snap = []
for f in [os.path.join(HIST, "000852_snapshot.parquet"),
os.path.join(HIST, "932000_snapshot.parquet")]:
if os.path.exists(f):
d = pd.read_parquet(f)[["index_code", "code", "code_name"]]
d["in_current"] = True
d["was_removed"] = False
d["source"] = "akshare_csindex"
snap.append(d)
df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame(
columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
df_snap["code"] = df_snap["code"].apply(norm_code)
print(f"[中证 snapshot] rows={len(df_snap)}")
# 3. 中证 1000/2000 announce_union 全集(优先); 缺则回退 snapshot 当前
ANNOUNCE_IDX = ["000852", "932000"]
df_csi = _read_announce_union_aggregated(HIST, ANNOUNCE_IDX)
if df_csi.empty:
# 回退: 旧 snapshot-only 路径(announce_union 未生成时)
snap = []
for f in [os.path.join(HIST, "000852_snapshot.parquet"),
os.path.join(HIST, "932000_snapshot.parquet")]:
if os.path.exists(f):
d = pd.read_parquet(f)[["index_code", "code", "code_name"]]
d["in_current"] = True
d["was_removed"] = False
d["source"] = "akshare_csindex"
snap.append(d)
df_csi = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame(
columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
df_csi["code"] = df_csi["code"].apply(norm_code)
print(f"[中证 snapshot 回退] rows={len(df_csi)}")
else:
print(f"[中证 announce_union] rows={len(df_csi)}, "
f"indices={df_csi['index_code'].nunique()}")
# 合并 + 去重 (同 index+code+source)
all_df = pd.concat([pool, df_deep, df_snap], ignore_index=True)
all_df = pd.concat([pool, df_deep, df_csi], ignore_index=True)
all_df = all_df.drop_duplicates(["index_code", "code", "source"])
print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, "
f"{all_df['index_code'].nunique()} indices")
@@ -163,7 +225,11 @@ def migrate(db_path: str) -> None:
"SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
f"{STAGING} WHERE index_code='399001'").fetchone())
print("sample 000852:", c.execute(
"SELECT COUNT(*) FROM " f"{STAGING} WHERE index_code='000852'").fetchone())
"SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
f"{STAGING} WHERE index_code='000852'").fetchone())
print("sample 932000:", c.execute(
"SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
f"{STAGING} WHERE index_code='932000'").fetchone())
finally:
c.close()
print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)")
scripts/data_platform/parse_csindex_announce.py
+175 -82
View File
@@ -118,6 +118,33 @@ def _download(url: str, target: Path, timeout: int = 60) -> int:
# ======================== 1. 列表 ========================
def _fetch_paginated(payload_base: dict, label: str) -> List[dict]:
"""通用分页拉取(单线程 + sleep)"""
all_items = []
page = 1
while True:
payload = dict(payload_base)
payload["page"] = {"desc": "", "key": "", "page": page, "rows": 100}
try:
d = _post_json(LIST_URL, payload)
except Exception as e:
log.error(f"[{label}] page {page} err: {e}; retry once after 5s")
time.sleep(5)
d = _post_json(LIST_URL, payload)
if d.get("code") != "200":
log.error(f"[{label}] page {page} API err: {d}")
break
items = d.get("data") or []
all_items.extend(items)
total = d.get("total") or 0
log.info(f" [{label}] page {page}: +{len(items)} (cum={len(all_items)}/{total})")
if not items or len(all_items) >= total:
break
page += 1
time.sleep(1.0)
return all_items
def fetch_all_notices(cache_path: Path, force: bool = False) -> List[dict]:
"""拉全量公告列表 (cached)"""
if cache_path.exists() and not force:
@@ -126,31 +153,11 @@ def fetch_all_notices(cache_path: Path, force: bool = False) -> List[dict]:
return json.load(f)
log.info(f"分页拉取全量公告: {LIST_URL}")
all_items = []
page = 1
while True:
payload = {
"lang": "cn", "classlist": [], "indexlist": [],
"page": {"desc": "", "key": "", "page": page, "rows": 100},
"related_topics": [], "typelist": [],
}
try:
d = _post_json(LIST_URL, payload)
except Exception as e:
log.error(f"page {page} err: {e}; retry once after 5s")
time.sleep(5)
d = _post_json(LIST_URL, payload)
if d.get("code") != "200":
log.error(f"page {page} API err: {d}")
break
items = d.get("data") or []
all_items.extend(items)
total = d.get("total") or 0
log.info(f" page {page}: +{len(items)} (cum={len(all_items)}/{total})")
if not items or len(all_items) >= total:
break
page += 1
time.sleep(1.0)
all_items = _fetch_paginated(
{"lang": "cn", "classlist": [], "indexlist": [],
"related_topics": [], "typelist": []},
label="all",
)
cache_path.parent.mkdir(parents=True, exist_ok=True)
with open(cache_path, "w", encoding="utf-8") as f:
@@ -159,6 +166,30 @@ def fetch_all_notices(cache_path: Path, force: bool = False) -> List[dict]:
return all_items
def fetch_notices_by_index(index_code: str, cache_path: Path, force: bool = False) -> List[dict]:
"""按 indexCode 精准分页拉公告列表 (cached)
实证: indexCode='000852' 返 96 条(44 调样, 回溯到 2007), 远优于 title 过滤的 28 条。
"""
if cache_path.exists() and not force:
log.info(f"使用缓存 indexCode 列表: {cache_path}")
with open(cache_path, encoding="utf-8") as f:
return json.load(f)
log.info(f"分页拉取 indexCode={index_code} 公告: {LIST_URL}")
items = _fetch_paginated(
{"lang": "cn", "classlist": [], "indexlist": [], "indexCode": index_code,
"related_topics": [], "typelist": []},
label=f"idx={index_code}",
)
cache_path.parent.mkdir(parents=True, exist_ok=True)
with open(cache_path, "w", encoding="utf-8") as f:
json.dump(items, f, ensure_ascii=False)
log.info(f"缓存 indexCode 列表: {cache_path} (total={len(items)})")
return items
def filter_csi1000_notices(notices: List[dict]) -> List[dict]:
"""筛 CSI 1000 调整公告 (theme=指数调样 + title 含 中证1000)"""
out = []
@@ -182,6 +213,29 @@ def filter_csi1000_notices(notices: List[dict]) -> List[dict]:
return out
def filter_adjustment_notices(notices: List[dict], keyword: str = "") -> List[dict]:
"""通用调整公告筛选: theme=指数调样 + (可选) title 含 keyword
用于 indexCode 拉取的列表(已按指数过滤,无需 title 匹配)。
"""
out = []
seen_ids = set()
for x in notices:
if x.get("theme") != "指数调样":
continue
title = x.get("title") or ""
if keyword and keyword not in title:
continue
if "不实施" in title:
continue
if x["id"] in seen_ids:
continue
seen_ids.add(x["id"])
out.append(x)
out.sort(key=lambda x: x.get("publishDate", ""))
return out
# ======================== 2. 详情 + 下载 ========================
def fetch_detail(nid: int, cache_dir: Path) -> dict:
"""详情 (cached by id)"""
@@ -429,11 +483,80 @@ def process_notice(nid: int, detail_cache: Path, file_cache: Path,
return records
def fetch_akshare_current(index_code: str) -> List[dict]:
"""拉 akshare 当前快照 -> records(adjust_type='current')
akshare 返列: 日期/指数代码/指数名称/指数英文名称/成分券代码/成分券名称/...
成分券代码 = 倒数第 2 不是 iloc[0](=日期), 用 header 名定位稳健。
"""
import akshare as ak
log.info(f"拉 akshare 当前快照 index={index_code}...")
df = ak.index_stock_cons_csindex(symbol=index_code)
# header 定位
cols = list(df.columns)
code_col = next((c for c in cols if "成分券代码" in str(c) or "股票代码" in str(c) or "证券代码" in str(c)), None)
name_col = next((c for c in cols if "成分券名称" in str(c) or "股票名称" in str(c) or "证券名称" in str(c)), None)
if code_col is None:
# 兜底: 第 5 列(成分券代码 实证位置)
code_col = cols[4] if len(cols) > 4 else cols[0]
if name_col is None and len(cols) > 5:
name_col = cols[5]
out = []
for _, row in df.iterrows():
raw = row[code_col]
s = str(raw).strip()
if not s or not s.isdigit():
continue
code = s.zfill(6)
if len(code) != 6:
continue
name = str(row[name_col]).strip() if name_col else ""
out.append({
"updateDate": "current", "index_code": index_code,
"code": code, "code_name": name,
"adjust_type": "current", "notice_id": 0, "source": "akshare.index_stock_cons_csindex",
})
log.info(f" akshare current {index_code}: {len(out)} stocks")
return out
def parse_launch_xlsx(path: Path, index_code: str, publish_date: str, notice_id: int) -> List[dict]:
"""解析 launch xlsx (单 sheet, header 6 列: 指数代码/指数简称/指数英文简称/证券代码/证券中文简称/证券英文名称)
用 header 定位"证券代码"和"证券中文简称"列, 不依赖固定 col index。
"""
wb = openpyxl.load_workbook(path, data_only=True)
out = []
for sn in wb.sheetnames:
ws = wb[sn]
rows = list(ws.iter_rows(values_only=True))
if not rows:
continue
header = rows[0]
code_idx = next((i for i, h in enumerate(header) if h and "证券代码" in str(h)), 3)
name_idx = next((i for i, h in enumerate(header) if h and "证券中文简称" in str(h)), 4)
for row in rows[1:]:
if not row:
continue
code = _norm_code(row[code_idx] if len(row) > code_idx else None)
if not code:
continue
name = str(row[name_idx]).strip() if len(row) > name_idx and row[name_idx] else ""
out.append({
"updateDate": publish_date, "index_code": index_code,
"code": code, "code_name": name,
"adjust_type": "initial", "notice_id": notice_id, "source": path.name,
})
break # 只用第一个 sheet
return out
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--out-dir", default=str(OUT_DIR))
parser.add_argument("--cache-dir", default=str(CACHE_DIR))
parser.add_argument("--refresh-list", action="store_true", help="强制重新拉全量公告列表")
parser.add_argument("--refresh-index", action="store_true", help="强制重新拉 indexCode 公告列表")
parser.add_argument("--only", choices=["1000", "2000", "both"], default="both")
args = parser.parse_args()
@@ -443,17 +566,29 @@ def main():
detail_cache = cache_dir / "detail"
file_cache = cache_dir / "files"
# 1. 拉/缓存全量列表
# 1. 拉/缓存全量列表(兜底) + indexCode 精准列表(主源)
notices = fetch_all_notices(cache_dir / "all_notices.json", force=args.refresh_list)
notices_000852 = fetch_notices_by_index(
"000852", cache_dir / "notices_000852.json", force=args.refresh_index,
)
# ========== CSI 1000 ==========
if args.only in ("1000", "both"):
log.info("\n========== CSI 1000 (000852) ==========")
# 过滤 + 附加已知 id
filtered = filter_csi1000_notices(notices)
log.info(f"filtered CSI 1000 adjustment notices: {len(filtered)}")
# 主源: indexCode=000852 精准拉, theme=指数调样
adj_idx = filter_adjustment_notices(notices_000852)
log.info(f"indexCode=000852 调整公告: {len(adj_idx)}")
# 兜底: 旧 title 过滤(防 indexCode 接口变动)
filtered_old = filter_csi1000_notices(notices)
log.info(f"title 过滤兜底: {len(filtered_old)}")
all_ids = sorted({x["id"] for x in filtered} | set(CSI1000_REGULAR_IDS) | set(CSI1000_TEMP_IDS))
all_ids = (
{x["id"] for x in adj_idx}
| {x["id"] for x in filtered_old}
| set(CSI1000_REGULAR_IDS)
| set(CSI1000_TEMP_IDS)
)
all_ids = sorted(all_ids)
log.info(f"total CSI 1000 notice ids to process: {len(all_ids)}")
records_1000 = []
@@ -464,24 +599,11 @@ def main():
except Exception as e:
log.error(f" {nid} FAILED: {e}")
# 加 initial 集合 (launch 2014-09 -> 2014-09-18 id=2998, 无初始样本 xlsx, 跳过)
# 加 current (akshare 现快照)
# 加 current (akshare 现快照, header 定位列, 修 iloc[0] bug)
try:
import akshare as ak
log.info("拉 akshare 中证1000 当前快照...")
df = ak.index_stock_cons_csindex(symbol="000852")
for _, row in df.iterrows():
code = str(row.iloc[0]).zfill(6) if str(row.iloc[0]).isdigit() else str(row.iloc[0])
if not code.isdigit() or len(code) != 6:
continue
records_1000.append({
"updateDate": "current", "index_code": "000852",
"code": code, "code_name": str(row.iloc[1]) if df.shape[1] > 1 else "",
"adjust_type": "current", "notice_id": 0, "source": "akshare.index_stock_cons_csindex",
})
log.info(f" akshare current: {sum(1 for r in records_1000 if r['adjust_type']=='current')} stocks")
records_1000.extend(fetch_akshare_current("000852"))
except Exception as e:
log.warning(f"akshare 当前快照拉取失败 (非致命): {e}")
log.warning(f"akshare 000852 当前快照失败 (非致命): {e}")
df_1000 = pd.DataFrame(records_1000, columns=[
"updateDate", "index_code", "code", "code_name", "adjust_type", "notice_id", "source"])
@@ -504,7 +626,7 @@ def main():
log.warning(" 仅可从 2023-08-10 launch xlsx 获取初始 2000 只样本")
records_2000 = []
# launch xlsx (initial)
# launch xlsx (initial) — header 定位列, 修 row[0] bug
try:
detail = fetch_detail(CSI2000_LAUNCH_ID, detail_cache)
if detail.get("data"):
@@ -515,45 +637,16 @@ def main():
for f in files:
if f.suffix.lower() != ".xlsx":
continue
wb = openpyxl.load_workbook(f, data_only=True)
for sn in wb.sheetnames:
ws = wb[sn]
rows = list(ws.iter_rows(values_only=True))
if not rows:
continue
header = rows[0]
# 找 code 列 (通常第 1 列 或 命名"证券代码"/"成分券代码")
for row in rows[1:]:
if not row:
continue
code = _norm_code(row[0] if len(row) > 0 else None)
if not code:
continue
name = str(row[1]).strip() if len(row) > 1 and row[1] else ""
records_2000.append({
"updateDate": publish_date, "index_code": "932000",
"code": code, "code_name": name,
"adjust_type": "initial", "notice_id": CSI2000_LAUNCH_ID, "source": f.name,
})
break # 只用第一个 xlsx
recs = parse_launch_xlsx(f, "932000", publish_date, CSI2000_LAUNCH_ID)
records_2000.extend(recs)
log.info(f" launch xlsx {f.name}: {len(recs)} initial stocks")
break
except Exception as e:
log.error(f"CSI 2000 launch xlsx ERR: {e}")
# current snapshot via akshare
# current snapshot via akshare (header 定位列)
try:
import akshare as ak
log.info("拉 akshare 中证2000 当前快照...")
df = ak.index_stock_cons_csindex(symbol="932000")
for _, row in df.iterrows():
code = str(row.iloc[0]).zfill(6) if str(row.iloc[0]).isdigit() else str(row.iloc[0])
if not code.isdigit() or len(code) != 6:
continue
records_2000.append({
"updateDate": "current", "index_code": "932000",
"code": code, "code_name": str(row.iloc[1]) if df.shape[1] > 1 else "",
"adjust_type": "current", "notice_id": 0, "source": "akshare.index_stock_cons_csindex",
})
log.info(f" akshare current: {sum(1 for r in records_2000 if r['adjust_type']=='current')} stocks")
records_2000.extend(fetch_akshare_current("932000"))
except Exception as e:
log.warning(f"akshare CSI 2000 当前快照失败: {e}")
tests/portfolio/test_migrate_announce_union.py
+191
View File
@@ -0,0 +1,191 @@
"""tests/portfolio/test_migrate_announce_union.py
TDD for migrate_constituent._read_announce_union_aggregated()
聚合逻辑(全集型):
- ever_codes = announce_union 所有 distinct code (add/remove/initial/current 任一)
- current_codes = announce_union 中 adjust_type='current' 的 code (∪ snapshot 兜底)
- in_current = code in current_codes
- was_removed = not in_current (曾经入选但当前不在)
对应 plan: docs/superpowers/plans/2026-07-23-csi1000-constituent-backfill.md Task 2
"""
import os
import sys
from pathlib import Path
import pandas as pd
import pytest
# 让测试能 import scripts/...
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "scripts" / "data_platform"))
# ---------- import target ----------
@pytest.fixture(scope="module")
def migrate_module():
"""import migrate_constituent 模块(不走 main, 只用函数)"""
import importlib.util
spec = importlib.util.spec_from_file_location(
"migrate_constituent_for_test",
ROOT / "scripts" / "data_platform" / "migrate_constituent.py",
)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
# ---------- 合成 announce_union ----------
def _make_announce_df(rows):
"""rows: list of (code, code_name, adjust_type)"""
return pd.DataFrame(
[{"updateDate": "2024-01-01", "index_code": "000852",
"code": c, "code_name": n, "adjust_type": t,
"notice_id": 1, "source": "synthetic.xlsx"} for c, n, t in rows]
)
# ============================================================
# Unit: 聚合逻辑(ever/in_current/was_removed)
# ============================================================
def test_aggregation_basic(migrate_module, tmp_path):
"""announce(add A,B + remove C) + current(A,B,D) -> ever={A,B,C,D},
in_current={A,B,D}, was_removed={C}
"""
# build announce_union parquet (含 current 行)
ann = _make_announce_df([
("000001", "A", "add"),
("000002", "B", "add"),
("000003", "C", "remove"),
("000001", "A", "current"),
("000002", "B", "current"),
("000004", "D", "current"),
])
hist = tmp_path / "hist"
hist.mkdir()
ann.to_parquet(hist / "000852_announce_union.parquet", index=False)
df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
assert len(df) == 4 # ever set 4 stocks
# index_code 全是 000852
assert set(df["index_code"]) == {"000852"}
# ever = {A,B,C,D}
assert set(df["code"]) == {"000001", "000002", "000003", "000004"}
# in_current = {A,B,D}, was_removed = {C}
in_cur = set(df[df["in_current"] == 1]["code"])
removed = set(df[df["was_removed"] == 1]["code"])
assert in_cur == {"000001", "000002", "000004"}
assert removed == {"000003"}
# 不存在既 in_current 又 was_removed 的行
assert ((df["in_current"] == 1) & (df["was_removed"] == 1)).sum() == 0
# 每行至少一个标记(全集型的硬条件)
assert ((df["in_current"] == 1) | (df["was_removed"] == 1)).sum() == 4
def test_aggregation_no_current_rows_uses_snapshot(migrate_module, tmp_path):
"""announce_union 只有 add/remove 时, 从 snapshot.parquet 兜底 current_codes"""
ann = _make_announce_df([
("000001", "A", "add"),
("000002", "B", "add"),
("000003", "C", "remove"),
])
snap = pd.DataFrame({
"index_code": ["000852", "000852"],
"code": ["000001", "000002"],
"code_name": ["A", "B"],
})
hist = tmp_path / "hist"
hist.mkdir()
ann.to_parquet(hist / "000852_announce_union.parquet", index=False)
snap.to_parquet(hist / "000852_snapshot.parquet", index=False)
df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
assert set(df["code"]) == {"000001", "000002", "000003"}
in_cur = set(df[df["in_current"] == 1]["code"])
assert in_cur == {"000001", "000002"} # snapshot 兜底
removed = set(df[df["was_removed"] == 1]["code"])
assert removed == {"000003"}
def test_aggregation_idempotent(migrate_module, tmp_path):
"""跑两次结果一致(幂等)"""
ann = _make_announce_df([
("000001", "A", "add"),
("000002", "B", "remove"),
("000001", "A", "current"),
])
hist = tmp_path / "hist"
hist.mkdir()
ann.to_parquet(hist / "000852_announce_union.parquet", index=False)
df1 = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
df2 = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
pd.testing.assert_frame_equal(
df1.reset_index(drop=True), df2.reset_index(drop=True),
check_like=True,
)
def test_aggregation_missing_file_returns_empty(migrate_module, tmp_path):
"""文件不存在 -> 空 DataFrame(不崩)"""
hist = tmp_path / "hist"
hist.mkdir()
df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
assert df.empty
assert list(df.columns) == [
"index_code", "code", "code_name", "in_current", "was_removed", "source"]
def test_aggregation_prefers_snapshot_name(migrate_module, tmp_path):
"""code_name 优先 snapshot 当前名(announce 的历史名可能过时)"""
ann = _make_announce_df([
("000001", "旧名", "add"),
("000001", "旧名", "current"),
])
snap = pd.DataFrame({
"index_code": ["000852"],
"code": ["000001"],
"code_name": ["新名"],
})
hist = tmp_path / "hist"
hist.mkdir()
ann.to_parquet(hist / "000852_announce_union.parquet", index=False)
snap.to_parquet(hist / "000852_snapshot.parquet", index=False)
df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"])
assert df.iloc[0]["code_name"] == "新名"
# ============================================================
# Integration: 真实 announce_union.parquet (Mac parse 产出)
# ============================================================
MAC_HIST = ROOT / "data" / "index_const_hist"
@pytest.mark.integration
def test_real_000852_distinct_greater_1000(migrate_module):
"""000852 治偏差证据: distinct > 1000(plan 验证标准)"""
if not (MAC_HIST / "000852_announce_union.parquet").exists():
pytest.skip("000852_announce_union.parquet 未生成, 先跑 parse_csindex_announce.py")
df = migrate_module._read_announce_union_aggregated(str(MAC_HIST), ["000852"])
n_distinct = df["code"].nunique()
n_removed = (df["was_removed"] == 1).sum()
assert n_distinct > 1000, f"000852 distinct={n_distinct}, 期望 >1000(治偏差)"
assert n_removed > 0, f"000852 was_removed={n_removed}, 期望 >0"
print(f"\n000852: distinct={n_distinct}, in_current={int((df['in_current']==1).sum())}, "
f"was_removed={n_removed}")
@pytest.mark.integration
def test_real_932000_distinct_approx_2000(migrate_module):
"""932000 launch xlsx bug 修复: distinct ≈ 2000+(launch ∪ current)"""
if not (MAC_HIST / "932000_announce_union.parquet").exists():
pytest.skip("932000_announce_union.parquet 未生成, 先跑 parse_csindex_announce.py")
df = migrate_module._read_announce_union_aggregated(str(MAC_HIST), ["932000"])
n_distinct = df["code"].nunique()
# launch 修复后 distinct 不再=1, 应在 2000+(launch ∪ current 重叠后 2684 实证)
assert n_distinct >= 2000, f"932000 distinct={n_distinct}, 期望 ≥2000(launch 修复)"
print(f"\n932000: distinct={n_distinct}, in_current={int((df['in_current']==1).sum())}, "
f"was_removed={int((df['was_removed']==1).sum())}")