claude_dev
6e719de0c6
E13 前置: adapter 逐文件读全市场不可行,先聚合成单 parquet (file_code/period/hold_pct/n_holders;null占比按0,坏文件名跳过,幂等--force)。 NAS 真跑已验证: 449.7s/108,610文件 → 97,836 行(skipped=4370 empty=6404)。
102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
# tests/factor/test_preaggregate_top_holders.py
|
||
"""top_holders 预聚合一次性脚本: 108,610 小文件 → 单 parquet 聚合产物.
|
||
|
||
契约(P1-B 任务书):
|
||
- 输出 (file_code, period, hold_pct, n_holders),hold_pct = 单文件 10 行
|
||
「占总流通股本持股比例」(%)求和;null 占比按 0
|
||
- 幂等: dst 存在默认跳过(--force 重算);绝不写 static 树
|
||
- 增减列混合类型(「不变」字符串)不参与求和,读入 coerce 不炸
|
||
"""
|
||
import os
|
||
import sys
|
||
from datetime import date
|
||
|
||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
|
||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "vnpy_v4.4.0")))
|
||
|
||
import polars as pl
|
||
import pytest
|
||
|
||
from scripts.factor_research.preaggregate_top_holders import aggregate, main
|
||
|
||
|
||
def _write_file(root: str, file_code: str, period: str, ratios, with_null=False):
|
||
os.makedirs(root, exist_ok=True)
|
||
rows = []
|
||
for r, ratio in enumerate(ratios):
|
||
rows.append({
|
||
"名次": r + 1, "股东名称": f"股东{r}", "股东性质": "基金",
|
||
"股份类型": "流通A股", "持股数": 1000 + r,
|
||
"占总流通股本持股比例": ratio,
|
||
"增减": "不变" if r % 2 else f"{1000 + r}", # 混合类型
|
||
"变动比率": 0.0,
|
||
})
|
||
if with_null:
|
||
rows.append({"名次": 11, "股东名称": "股东10", "股东性质": "基金",
|
||
"股份类型": "流通A股", "持股数": 1,
|
||
"占总流通股本持股比例": None, "增减": "不变", "变动比率": 0.0})
|
||
pl.DataFrame(rows).write_parquet(
|
||
os.path.join(root, f"{file_code}_{period}_top_holders.parquet"))
|
||
|
||
|
||
@pytest.fixture()
|
||
def th_tree(tmp_path):
|
||
root = str(tmp_path / "top_holders")
|
||
_write_file(root, "600000.SH", "20221231", [6.0] * 10)
|
||
# null 占比行(9 行有值 + 1 行 null → sum=45,不因 null 炸/漏)
|
||
_write_file(root, "000001.SZ", "20221231", [5.0] * 9, with_null=True)
|
||
_write_file(root, "600000.SH", "20230331", [6.5] * 10)
|
||
return root
|
||
|
||
|
||
def test_aggregate_sums_per_code_period(th_tree, tmp_path):
|
||
dst = str(tmp_path / "factor_cache" / "agg.parquet")
|
||
df = aggregate(th_tree, dst)
|
||
assert set(df.columns) == {"file_code", "period", "hold_pct", "n_holders"}
|
||
got = {(r["file_code"], r["period"]): r["hold_pct"] for r in df.iter_rows(named=True)}
|
||
assert got == {
|
||
("600000.SH", date(2022, 12, 31)): 60.0,
|
||
("600000.SH", date(2023, 3, 31)): 65.0,
|
||
("000001.SZ", date(2022, 12, 31)): 45.0, # null 行按 0,9×5.0
|
||
}
|
||
# n_holders = 原始行数(null 行也计入)
|
||
n = {(r["file_code"], r["period"]): r["n_holders"] for r in df.iter_rows(named=True)}
|
||
assert n[("000001.SZ", date(2022, 12, 31))] == 10
|
||
|
||
|
||
def test_aggregate_idempotent_and_force(th_tree, tmp_path, capsys):
|
||
dst = str(tmp_path / "factor_cache" / "agg.parquet")
|
||
aggregate(th_tree, dst)
|
||
first_mtime = os.path.getmtime(dst)
|
||
# 已存在 → 跳过重算(mtime 不变)
|
||
df2 = aggregate(th_tree, dst)
|
||
assert os.path.getmtime(dst) == first_mtime
|
||
assert df2.height == 3
|
||
out = capsys.readouterr().out
|
||
assert "跳过" in out
|
||
# --force 重算(mtime 更新;新文件落盘触发)
|
||
import time
|
||
time.sleep(0.01)
|
||
_write_file(th_tree, "600000.SH", "20230630", [6.2] * 10)
|
||
aggregate(th_tree, dst, force=True)
|
||
assert os.path.getmtime(dst) > first_mtime
|
||
df3 = pl.read_parquet(dst)
|
||
assert df3.height == 4
|
||
|
||
|
||
def test_aggregate_skips_bad_filenames(th_tree, tmp_path):
|
||
# 非契约文件名(北交/杂名)跳过不炸
|
||
_write_file(th_tree, "920001.BJ", "20221231", [1.0] * 10)
|
||
(tmp_path / "top_holders" / "notes.txt").write_text("junk")
|
||
dst = str(tmp_path / "agg2.parquet")
|
||
df = aggregate(th_tree, dst)
|
||
assert df.filter(pl.col("file_code") == "920001.BJ").height == 0
|
||
|
||
|
||
def test_main_cli(th_tree, tmp_path, capsys):
|
||
dst = str(tmp_path / "cli" / "agg.parquet")
|
||
main(["--src", th_tree, "--dst", dst])
|
||
df = pl.read_parquet(dst)
|
||
assert df.height == 3
|
||
assert "rows" in capsys.readouterr().out
|