fix(data): 中证1000/2000 历史成份股补全(parse+migrate 聚合)
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:
@@ -77,6 +77,61 @@ def _read_baostock_constituent(c: sqlite3.Connection) -> pd.DataFrame:
|
|||||||
return df
|
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:
|
def migrate(db_path: str) -> None:
|
||||||
"""(幂等)重建 constituent_unified_staging。
|
"""(幂等)重建 constituent_unified_staging。
|
||||||
|
|
||||||
@@ -115,23 +170,30 @@ def migrate(db_path: str) -> None:
|
|||||||
df_deep["code"] = df_deep["code"].apply(norm_code)
|
df_deep["code"] = df_deep["code"].apply(norm_code)
|
||||||
print(f"[深证 union] rows={len(df_deep)}")
|
print(f"[深证 union] rows={len(df_deep)}")
|
||||||
|
|
||||||
# 3. 中证 snapshot
|
# 3. 中证 1000/2000 announce_union 全集(优先); 缺则回退 snapshot 当前
|
||||||
snap = []
|
ANNOUNCE_IDX = ["000852", "932000"]
|
||||||
for f in [os.path.join(HIST, "000852_snapshot.parquet"),
|
df_csi = _read_announce_union_aggregated(HIST, ANNOUNCE_IDX)
|
||||||
os.path.join(HIST, "932000_snapshot.parquet")]:
|
if df_csi.empty:
|
||||||
if os.path.exists(f):
|
# 回退: 旧 snapshot-only 路径(announce_union 未生成时)
|
||||||
d = pd.read_parquet(f)[["index_code", "code", "code_name"]]
|
snap = []
|
||||||
d["in_current"] = True
|
for f in [os.path.join(HIST, "000852_snapshot.parquet"),
|
||||||
d["was_removed"] = False
|
os.path.join(HIST, "932000_snapshot.parquet")]:
|
||||||
d["source"] = "akshare_csindex"
|
if os.path.exists(f):
|
||||||
snap.append(d)
|
d = pd.read_parquet(f)[["index_code", "code", "code_name"]]
|
||||||
df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame(
|
d["in_current"] = True
|
||||||
columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
|
d["was_removed"] = False
|
||||||
df_snap["code"] = df_snap["code"].apply(norm_code)
|
d["source"] = "akshare_csindex"
|
||||||
print(f"[中证 snapshot] rows={len(df_snap)}")
|
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)
|
# 合并 + 去重 (同 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"])
|
all_df = all_df.drop_duplicates(["index_code", "code", "source"])
|
||||||
print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, "
|
print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, "
|
||||||
f"{all_df['index_code'].nunique()} indices")
|
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 "
|
"SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
|
||||||
f"{STAGING} WHERE index_code='399001'").fetchone())
|
f"{STAGING} WHERE index_code='399001'").fetchone())
|
||||||
print("sample 000852:", c.execute(
|
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:
|
finally:
|
||||||
c.close()
|
c.close()
|
||||||
print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)")
|
print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)")
|
||||||
|
|||||||
@@ -118,6 +118,33 @@ def _download(url: str, target: Path, timeout: int = 60) -> int:
|
|||||||
|
|
||||||
|
|
||||||
# ======================== 1. 列表 ========================
|
# ======================== 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]:
|
def fetch_all_notices(cache_path: Path, force: bool = False) -> List[dict]:
|
||||||
"""拉全量公告列表 (cached)"""
|
"""拉全量公告列表 (cached)"""
|
||||||
if cache_path.exists() and not force:
|
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)
|
return json.load(f)
|
||||||
|
|
||||||
log.info(f"分页拉取全量公告: {LIST_URL}")
|
log.info(f"分页拉取全量公告: {LIST_URL}")
|
||||||
all_items = []
|
all_items = _fetch_paginated(
|
||||||
page = 1
|
{"lang": "cn", "classlist": [], "indexlist": [],
|
||||||
while True:
|
"related_topics": [], "typelist": []},
|
||||||
payload = {
|
label="all",
|
||||||
"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)
|
|
||||||
|
|
||||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
with open(cache_path, "w", encoding="utf-8") as f:
|
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
|
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]:
|
def filter_csi1000_notices(notices: List[dict]) -> List[dict]:
|
||||||
"""筛 CSI 1000 调整公告 (theme=指数调样 + title 含 中证1000)"""
|
"""筛 CSI 1000 调整公告 (theme=指数调样 + title 含 中证1000)"""
|
||||||
out = []
|
out = []
|
||||||
@@ -182,6 +213,29 @@ def filter_csi1000_notices(notices: List[dict]) -> List[dict]:
|
|||||||
return out
|
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. 详情 + 下载 ========================
|
# ======================== 2. 详情 + 下载 ========================
|
||||||
def fetch_detail(nid: int, cache_dir: Path) -> dict:
|
def fetch_detail(nid: int, cache_dir: Path) -> dict:
|
||||||
"""详情 (cached by id)"""
|
"""详情 (cached by id)"""
|
||||||
@@ -429,11 +483,80 @@ def process_notice(nid: int, detail_cache: Path, file_cache: Path,
|
|||||||
return records
|
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():
|
def main():
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument("--out-dir", default=str(OUT_DIR))
|
parser.add_argument("--out-dir", default=str(OUT_DIR))
|
||||||
parser.add_argument("--cache-dir", default=str(CACHE_DIR))
|
parser.add_argument("--cache-dir", default=str(CACHE_DIR))
|
||||||
parser.add_argument("--refresh-list", action="store_true", help="强制重新拉全量公告列表")
|
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")
|
parser.add_argument("--only", choices=["1000", "2000", "both"], default="both")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
@@ -443,17 +566,29 @@ def main():
|
|||||||
detail_cache = cache_dir / "detail"
|
detail_cache = cache_dir / "detail"
|
||||||
file_cache = cache_dir / "files"
|
file_cache = cache_dir / "files"
|
||||||
|
|
||||||
# 1. 拉/缓存全量列表
|
# 1. 拉/缓存全量列表(兜底) + indexCode 精准列表(主源)
|
||||||
notices = fetch_all_notices(cache_dir / "all_notices.json", force=args.refresh_list)
|
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 ==========
|
# ========== CSI 1000 ==========
|
||||||
if args.only in ("1000", "both"):
|
if args.only in ("1000", "both"):
|
||||||
log.info("\n========== CSI 1000 (000852) ==========")
|
log.info("\n========== CSI 1000 (000852) ==========")
|
||||||
# 过滤 + 附加已知 id
|
# 主源: indexCode=000852 精准拉, theme=指数调样
|
||||||
filtered = filter_csi1000_notices(notices)
|
adj_idx = filter_adjustment_notices(notices_000852)
|
||||||
log.info(f"filtered CSI 1000 adjustment notices: {len(filtered)}")
|
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)}")
|
log.info(f"total CSI 1000 notice ids to process: {len(all_ids)}")
|
||||||
|
|
||||||
records_1000 = []
|
records_1000 = []
|
||||||
@@ -464,24 +599,11 @@ def main():
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.error(f" {nid} FAILED: {e}")
|
log.error(f" {nid} FAILED: {e}")
|
||||||
|
|
||||||
# 加 initial 集合 (launch 2014-09 -> 2014-09-18 id=2998, 无初始样本 xlsx, 跳过)
|
# 加 current (akshare 现快照, header 定位列, 修 iloc[0] bug)
|
||||||
# 加 current (akshare 现快照)
|
|
||||||
try:
|
try:
|
||||||
import akshare as ak
|
records_1000.extend(fetch_akshare_current("000852"))
|
||||||
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")
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning(f"akshare 当前快照拉取失败 (非致命): {e}")
|
log.warning(f"akshare 000852 当前快照失败 (非致命): {e}")
|
||||||
|
|
||||||
df_1000 = pd.DataFrame(records_1000, columns=[
|
df_1000 = pd.DataFrame(records_1000, columns=[
|
||||||
"updateDate", "index_code", "code", "code_name", "adjust_type", "notice_id", "source"])
|
"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 只样本")
|
log.warning(" 仅可从 2023-08-10 launch xlsx 获取初始 2000 只样本")
|
||||||
records_2000 = []
|
records_2000 = []
|
||||||
|
|
||||||
# launch xlsx (initial)
|
# launch xlsx (initial) — header 定位列, 修 row[0] bug
|
||||||
try:
|
try:
|
||||||
detail = fetch_detail(CSI2000_LAUNCH_ID, detail_cache)
|
detail = fetch_detail(CSI2000_LAUNCH_ID, detail_cache)
|
||||||
if detail.get("data"):
|
if detail.get("data"):
|
||||||
@@ -515,45 +637,16 @@ def main():
|
|||||||
for f in files:
|
for f in files:
|
||||||
if f.suffix.lower() != ".xlsx":
|
if f.suffix.lower() != ".xlsx":
|
||||||
continue
|
continue
|
||||||
wb = openpyxl.load_workbook(f, data_only=True)
|
recs = parse_launch_xlsx(f, "932000", publish_date, CSI2000_LAUNCH_ID)
|
||||||
for sn in wb.sheetnames:
|
records_2000.extend(recs)
|
||||||
ws = wb[sn]
|
log.info(f" launch xlsx {f.name}: {len(recs)} initial stocks")
|
||||||
rows = list(ws.iter_rows(values_only=True))
|
break
|
||||||
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
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.error(f"CSI 2000 launch xlsx ERR: {e}")
|
log.error(f"CSI 2000 launch xlsx ERR: {e}")
|
||||||
|
|
||||||
# current snapshot via akshare
|
# current snapshot via akshare (header 定位列)
|
||||||
try:
|
try:
|
||||||
import akshare as ak
|
records_2000.extend(fetch_akshare_current("932000"))
|
||||||
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")
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning(f"akshare CSI 2000 当前快照失败: {e}")
|
log.warning(f"akshare CSI 2000 当前快照失败: {e}")
|
||||||
|
|
||||||
|
|||||||
@@ -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())}")
|
||||||
Reference in New Issue
Block a user