diff --git a/scripts/data_platform/migrate_constituent.py b/scripts/data_platform/migrate_constituent.py index da8dc29..821c3dd 100644 --- a/scripts/data_platform/migrate_constituent.py +++ b/scripts/data_platform/migrate_constituent.py @@ -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 + (∪ _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 后单独合并)") diff --git a/scripts/data_platform/parse_csindex_announce.py b/scripts/data_platform/parse_csindex_announce.py index 31e322c..c65ebff 100644 --- a/scripts/data_platform/parse_csindex_announce.py +++ b/scripts/data_platform/parse_csindex_announce.py @@ -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}") diff --git a/tests/portfolio/test_migrate_announce_union.py b/tests/portfolio/test_migrate_announce_union.py new file mode 100644 index 0000000..be8e47d --- /dev/null +++ b/tests/portfolio/test_migrate_announce_union.py @@ -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())}")