claude_dev
b237c2d7e7
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
192 lines
7.5 KiB
Python
192 lines
7.5 KiB
Python
"""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())}")
|