claude_dev
774170ec05
采集层(多源各下): - baostock: 日线全字段全量(baostock_daily_fullmarket) + 15min全市场 + 静态(基础/复权/分红/季频/三表) + 成份股 - akshare: 静态(估值/龙虎榜/大宗/融资融券/北向/指数成分/行业/股本/解禁/业绩预告) - xtdata(miniQMT): build_daily_from_xtdata + daily_update_xtdata 数据补全 P0: - ETF全市场: universe 扩展 沪深A股∪ETF∪基金(7414), dividend_type='front' 前复权 - 历史成份股(治幸存者偏差): index_const_hist_download 深证/国证 adjust_cni 4指数 + 中证1000/2000快照 + 新浪交叉校验 - 退市K线: baostock_delisted_download + import_delisted_to_db(实证 Day1 fetch_all_stocks 已含退市) 灌库: - import_baostock_to_db: daily_baostock_full(5537股/1826万行,18字段)+ bs_index_constituent + bs_adjust_factor - INSERT OR REPLACE 幂等, WAL+busy_timeout, dbbardata 不碰 每日增量 #7(用户决策A: VPS直跑): - daily_update_static: login探针防黑名单graceful skip + LOOKBACK7 + query_stock_basic含退市 + INSERT OR REPLACE + QUERY_COUNT守48000/天 设计文档: spec(13节三层融合) + P0 plan + 数据gap设计
497 lines
16 KiB
Python
Executable File
497 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""BaoStock 成分股历史下载脚本 (只产 parquet)
|
||
|
||
目标:
|
||
在 Mac 本机用 baostock 拉取 HS300/ZZ500/SZ50 历史成分股快照,
|
||
按历史日期循环 (每周一一个快照, 从 2006-01 至今)。
|
||
输出 parquet 到 staging 目录。
|
||
|
||
硬约束 (踩过的坑):
|
||
1. 单进程单登录, 严禁并发 (baostock 并发会拉黑封 IP 6-24h)
|
||
2. 直连不走代理 (脚本开头 unset proxy)
|
||
3. 每次 query 后 sleep 0.4s 限速
|
||
4. staging parquet, 绝不直接写主库
|
||
|
||
用法:
|
||
python baostock_constituent_download.py # 全指数 2006-01~至今
|
||
python baostock_constituent_download.py --limit 5 # 测试: 前 5 个快照
|
||
python baostock_constituent_download.py --indices hs300 # 只跑指定指数
|
||
python baostock_constituent_download.py --start 2020-01-01 # 指定起始日期
|
||
|
||
退出码: 0=完成, 1=致命错误, 2=断路器触发 (可重试)
|
||
"""
|
||
|
||
import argparse
|
||
import datetime
|
||
import json
|
||
import logging
|
||
import os
|
||
import socket
|
||
import sys
|
||
import time
|
||
from datetime import date, timedelta
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
# ======================== 硬约束: unset proxy + socket timeout ========================
|
||
for _k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
|
||
os.environ.pop(_k, None)
|
||
|
||
socket.setdefaulttimeout(30)
|
||
|
||
try:
|
||
sys.stdout.reconfigure(line_buffering=True)
|
||
except (AttributeError, ValueError):
|
||
pass
|
||
|
||
import baostock as bs # noqa: E402
|
||
import pandas as pd # noqa: E402
|
||
|
||
|
||
# ======================== 配置 ========================
|
||
|
||
DEFAULT_OUT_DIR = "/Users/chufeng/.openclaw/sanguo_projects/sanguo_vnpy_v2/data/constituent_baostock"
|
||
OUT_DIR = Path(os.environ.get("BS_CONSTITUENT_OUT_DIR", DEFAULT_OUT_DIR))
|
||
DEFAULT_LOG_DIR = OUT_DIR / "logs"
|
||
_log_dir_env = os.environ.get("BS_CONSTITUENT_LOG_DIR")
|
||
LOG_DIR = Path(_log_dir_env) if _log_dir_env else DEFAULT_LOG_DIR
|
||
|
||
START_DATE_DEFAULT = "2006-01-01" # 成分股历史起始
|
||
END_DATE_DEFAULT = "" # 默认今天
|
||
|
||
BS_INTERVAL = 0.4 # 每次间隔秒 (防封 IP)
|
||
BS_MAX_RETRIES = 3 # 单次下载重试次数
|
||
PROGRESS_LOG_EVERY = 50 # 每 N 个快照打一次进度日志
|
||
RELOGIN_EVERY = 100 # 每 N 个快照定期重登
|
||
CIRCUIT_BREAKER = 10 # 连续失败 N 次 → 断路退出
|
||
|
||
# 指数配置 (baostock code, 显示名)
|
||
INDICES_CONFIG = {
|
||
"hs300": {"code": "000300", "name": "HS300"},
|
||
"zz500": {"code": "000905", "name": "ZZ500"},
|
||
"sz50": {"code": "000016", "name": "SZ50"},
|
||
}
|
||
|
||
|
||
# ======================== 日志 ========================
|
||
|
||
def setup_logging() -> Tuple[logging.Logger, Path]:
|
||
"""配置 root logger: 同时写 stdout + 文件。"""
|
||
LOG_DIR.mkdir(parents=True, exist_ok=True)
|
||
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
log_file = LOG_DIR / f"baostock_constituent_{ts}.log"
|
||
|
||
fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
|
||
root = logging.getLogger()
|
||
root.setLevel(logging.INFO)
|
||
for h in list(root.handlers):
|
||
root.removeHandler(h)
|
||
|
||
sh = logging.StreamHandler(sys.stdout)
|
||
sh.setFormatter(fmt)
|
||
root.addHandler(sh)
|
||
|
||
fh = logging.FileHandler(log_file, encoding="utf-8")
|
||
fh.setFormatter(fmt)
|
||
root.addHandler(fh)
|
||
|
||
return logging.getLogger(__name__), log_file
|
||
|
||
|
||
logger, LOG_FILE = setup_logging()
|
||
|
||
|
||
# ======================== 工具函数 ========================
|
||
|
||
def normalize_date(s: str) -> str:
|
||
"""接受 YYYYMMDD 或 YYYY-MM-DD, 返回 YYYY-MM-DD; 空串 → 今天。"""
|
||
s = (s or "").strip()
|
||
if not s:
|
||
return datetime.date.today().strftime("%Y-%m-%d")
|
||
if "-" in s:
|
||
return s
|
||
if len(s) == 8 and s.isdigit():
|
||
return f"{s[:4]}-{s[4:6]}-{s[6:8]}"
|
||
raise ValueError(f"无效日期格式: {s}")
|
||
|
||
|
||
def get_monday_dates(start_date: str, end_date: str) -> List[str]:
|
||
"""生成每周一日期列表 (包含 start_date 所在周, 含 end_date)。
|
||
|
||
返回: YYYY-MM-DD 格式日期列表。
|
||
"""
|
||
start = datetime.datetime.strptime(start_date, "%Y-%m-%d").date()
|
||
end = datetime.datetime.strptime(end_date, "%Y-%m-%d").date()
|
||
|
||
# 找到 start_date 后的第一个周一
|
||
d = start
|
||
while d.weekday() != 0: # Monday=0
|
||
d += timedelta(days=1)
|
||
|
||
dates = []
|
||
while d <= end:
|
||
dates.append(d.strftime("%Y-%m-%d"))
|
||
d += timedelta(days=7) # 下一周
|
||
|
||
return dates
|
||
|
||
|
||
def baostock_code_to_parts(bs_code: str) -> Tuple[str, str]:
|
||
"""sh.600519 → (600519, 'SH'); sz.000001 → (000001, 'SZ')。"""
|
||
if "." not in bs_code:
|
||
raise ValueError(f"无效 baostock code: {bs_code}")
|
||
prefix, num = bs_code.split(".", 1)
|
||
if prefix == "sh":
|
||
return num, "SH"
|
||
if prefix == "sz":
|
||
return num, "SZ"
|
||
raise ValueError(f"未知 baostock 前缀: {bs_code}")
|
||
|
||
|
||
# ======================== baostock 登录 ========================
|
||
|
||
def _login_once() -> bool:
|
||
"""bs.login() — 已被 socket.setdefaulttimeout(30) 保护。"""
|
||
try:
|
||
lg = bs.login()
|
||
if lg.error_code == "0":
|
||
return True
|
||
logger.error("bs.login() 失败: code=%s msg=%s", lg.error_code, lg.error_msg)
|
||
return False
|
||
except (socket.timeout, TimeoutError) as e:
|
||
logger.error("bs.login() 超时: %s (baostock 疑似冷却)", e)
|
||
return False
|
||
except Exception as e:
|
||
logger.error("bs.login() 异常: %s", e)
|
||
return False
|
||
|
||
|
||
def _relogin() -> bool:
|
||
"""强制重登: logout + login。失败等 2s 再试 1 次。"""
|
||
try:
|
||
bs.logout()
|
||
except Exception:
|
||
pass
|
||
if _login_once():
|
||
return True
|
||
time.sleep(2)
|
||
try:
|
||
bs.logout()
|
||
except Exception:
|
||
pass
|
||
return _login_once()
|
||
|
||
|
||
# ======================== 数据下载 ========================
|
||
|
||
def fetch_constituent_stocks(index_name: str, date_str: str) -> pd.DataFrame:
|
||
"""拉取指定指数在指定日期的成分股。
|
||
|
||
返回 df 列: code (baostock格式, 如 sh.600519)
|
||
空数据返回空 df。
|
||
"""
|
||
index_info = INDICES_CONFIG.get(index_name)
|
||
if not index_info:
|
||
raise ValueError(f"未知指数: {index_name}")
|
||
|
||
index_code = index_info["code"]
|
||
|
||
if index_name == "hs300":
|
||
rs = bs.query_hs300_stocks(date=date_str)
|
||
elif index_name == "zz500":
|
||
rs = bs.query_zz500_stocks(date=date_str)
|
||
elif index_name == "sz50":
|
||
rs = bs.query_sz50_stocks(date=date_str)
|
||
else:
|
||
raise ValueError(f"未实现的指数: {index_name}")
|
||
|
||
if rs.error_code != "0":
|
||
raise RuntimeError(
|
||
f"query_{index_name}_stocks({date_str}) 错误: {rs.error_code} {rs.error_msg}"
|
||
)
|
||
|
||
if rs.error_code != "0":
|
||
raise RuntimeError(
|
||
f"query_{index_name}_stocks({date_str}) 错误: {rs.error_code} {rs.error_msg}"
|
||
)
|
||
|
||
# 获取字段列表
|
||
fields = list(rs.fields)
|
||
|
||
rows = []
|
||
while rs.next():
|
||
rows.append(rs.get_row_data())
|
||
|
||
if not rows:
|
||
# 合法空数据 (指数不存在于该日期)
|
||
return pd.DataFrame(columns=fields)
|
||
|
||
# 使用 baostock 实际返回的字段名
|
||
return pd.DataFrame(rows, columns=fields)
|
||
|
||
|
||
# ======================== 路径 / marker ========================
|
||
|
||
def parquet_path_for(index_name: str, date_str: str) -> Path:
|
||
"""hs300 + 2020-01-06 → OUT_DIR / hs300_2020-01-06.parquet"""
|
||
return OUT_DIR / f"{index_name}_{date_str}.parquet"
|
||
|
||
|
||
def marker_path_for(parquet_path: Path) -> Path:
|
||
"""parquet → .{stem}.baostock marker"""
|
||
return parquet_path.parent / f".{parquet_path.stem}.baostock"
|
||
|
||
|
||
def load_done_set(indices: List[str]) -> Dict[str, set]:
|
||
"""扫 OUT_DIR 所有 marker 构造已完成集合。
|
||
|
||
返回: {index_name: set(dates)} 字典。
|
||
"""
|
||
done: Dict[str, set] = {idx: set() for idx in indices}
|
||
if not OUT_DIR.exists():
|
||
return done
|
||
|
||
suffix = ".baostock"
|
||
for marker in OUT_DIR.glob(f".*{suffix}"):
|
||
name = marker.name
|
||
if not name.startswith(".") or not name.endswith(suffix):
|
||
continue
|
||
stem = name[1:-len(suffix)] # e.g. hs300_2020-01-06
|
||
|
||
# 解析: <index_name>_<date>
|
||
try:
|
||
parts = stem.split("_", 1)
|
||
if len(parts) != 2:
|
||
continue
|
||
index_name, date_str = parts
|
||
if index_name in done:
|
||
done[index_name].add(date_str)
|
||
except ValueError:
|
||
continue
|
||
return done
|
||
|
||
|
||
# ======================== 单次下载 ========================
|
||
|
||
def download_one_snapshot(
|
||
index_name: str,
|
||
date_str: str,
|
||
force: bool,
|
||
) -> Tuple[str, int]:
|
||
"""下载单个成分股快照 → 写 parquet + marker。
|
||
|
||
返回 (status, rows): status ∈ {'ok', 'skipped', 'failed', 'empty'}。
|
||
"""
|
||
parquet_path = parquet_path_for(index_name, date_str)
|
||
marker_path = marker_path_for(parquet_path)
|
||
|
||
if not force and marker_path.exists():
|
||
return "skipped", 0
|
||
|
||
# retry 循环
|
||
df: Optional[pd.DataFrame] = None
|
||
for attempt in range(BS_MAX_RETRIES):
|
||
try:
|
||
df = fetch_constituent_stocks(index_name, date_str)
|
||
break
|
||
except (socket.timeout, TimeoutError, OSError) as e:
|
||
logger.warning(
|
||
"%s %s socket 异常重试 %d/%d: %s — 强制重登",
|
||
index_name, date_str, attempt + 1, BS_MAX_RETRIES, e,
|
||
)
|
||
if not _relogin():
|
||
logger.error("重登失败, 放弃 %s %s", index_name, date_str)
|
||
df = None
|
||
break
|
||
except Exception as e:
|
||
logger.warning(
|
||
"%s %s 下载异常重试 %d/%d: %s — 强制重登",
|
||
index_name, date_str, attempt + 1, BS_MAX_RETRIES, e,
|
||
)
|
||
if not _relogin():
|
||
logger.error("重登失败, 放弃 %s %s", index_name, date_str)
|
||
df = None
|
||
break
|
||
|
||
if df is None:
|
||
return "failed", 0
|
||
|
||
if df.empty:
|
||
# 合法空数据 (指数不存在于该日期) — 也写 marker
|
||
try:
|
||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||
df.to_parquet(parquet_path, index=False)
|
||
marker_path.write_text(datetime.datetime.now().isoformat())
|
||
return "empty", 0
|
||
except Exception as e:
|
||
logger.error("写入 %s 失败: %s", parquet_path, e)
|
||
return "failed", 0
|
||
|
||
# 写 parquet + marker
|
||
try:
|
||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||
df.to_parquet(parquet_path, index=False)
|
||
marker_path.write_text(datetime.datetime.now().isoformat())
|
||
return "ok", len(df)
|
||
except Exception as e:
|
||
logger.error("写入 %s 失败: %s", parquet_path, e)
|
||
return "failed", 0
|
||
|
||
|
||
# ======================== 主流程 ========================
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(
|
||
description="BaoStock 成分股历史下载 (只产 parquet)",
|
||
)
|
||
p.add_argument(
|
||
"--start", default=START_DATE_DEFAULT,
|
||
help=f"起始日期 YYYYMMDD 或 YYYY-MM-DD, 默认 {START_DATE_DEFAULT}",
|
||
)
|
||
p.add_argument("--end", default=END_DATE_DEFAULT, help="结束日期, 默认今天")
|
||
p.add_argument(
|
||
"--indices",
|
||
default="hs300,zz500,sz50",
|
||
help="逗号分隔指数名称, 默认 hs300,zz500,sz50",
|
||
)
|
||
p.add_argument("--limit", type=int, default=0, help="限制处理快照数 (测试用)")
|
||
p.add_argument("--force", action="store_true", help="强制重下, 忽略 marker")
|
||
return p.parse_args()
|
||
|
||
|
||
def resolve_indices(s: str) -> List[str]:
|
||
"""解析 --indices 字符串 → 去重保序的指数列表。"""
|
||
parts = [p.strip() for p in s.split(",") if p.strip()]
|
||
valid = set(INDICES_CONFIG.keys())
|
||
bad = [p for p in parts if p not in valid]
|
||
if bad:
|
||
raise SystemExit(f"未知 --indices: {bad}, 可选 {list(valid)}")
|
||
seen: set = set()
|
||
out: List[str] = []
|
||
for p in parts:
|
||
if p not in seen:
|
||
out.append(p)
|
||
seen.add(p)
|
||
return out
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
start_date = normalize_date(args.start)
|
||
end_date = normalize_date(args.end)
|
||
indices = resolve_indices(args.indices)
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("BaoStock 成分股历史下载")
|
||
logger.info(" 输出目录: %s", OUT_DIR)
|
||
logger.info(" 日志文件: %s", LOG_FILE)
|
||
logger.info(" 日期范围: %s ~ %s", start_date, end_date)
|
||
logger.info(" 指数: %s", indices)
|
||
logger.info(" 频率: 每周一快照")
|
||
logger.info(" socket.setdefaulttimeout(30)")
|
||
logger.info(" 当前时间: %s", datetime.datetime.now().isoformat())
|
||
logger.info("=" * 60)
|
||
|
||
# 登录
|
||
if not _login_once():
|
||
logger.error("[FATAL] baostock 登录失败 (疑似冷却), 退出")
|
||
sys.exit(1)
|
||
logger.info("baostock 登录成功")
|
||
|
||
# 生成日期列表 (每周一)
|
||
monday_dates = get_monday_dates(start_date, end_date)
|
||
logger.info("快照日期列表: %d 个周一 (%s ~ %s)", len(monday_dates),
|
||
monday_dates[0] if monday_dates else "N/A",
|
||
monday_dates[-1] if monday_dates else "N/A")
|
||
|
||
# limit
|
||
if args.limit > 0:
|
||
monday_dates = monday_dates[:args.limit]
|
||
logger.info("limit=%d 截断", args.limit)
|
||
|
||
# 总快照数
|
||
total_snapshots = len(monday_dates) * len(indices)
|
||
logger.info("总快照数: %d (日期%d × 指数%d)", total_snapshots,
|
||
len(monday_dates), len(indices))
|
||
|
||
# 断点续传
|
||
done_set = load_done_set(indices)
|
||
if args.force:
|
||
todo_snapshots = [(idx, d) for idx in indices for d in monday_dates]
|
||
else:
|
||
todo_snapshots = [
|
||
(idx, d) for idx in indices for d in monday_dates
|
||
if d not in done_set.get(idx, set())
|
||
]
|
||
logger.info("待处理: %d (已完成 %d)", len(todo_snapshots),
|
||
total_snapshots - len(todo_snapshots))
|
||
|
||
# 主循环
|
||
stats = {"ok": 0, "skipped": 0, "empty": 0, "failed": 0, "rows": 0}
|
||
consec_fail = 0
|
||
circuit_triggered = False
|
||
t_start = time.time()
|
||
|
||
for i, (index_name, date_str) in enumerate(todo_snapshots):
|
||
# 定期重登
|
||
if i > 0 and i % RELOGIN_EVERY == 0:
|
||
logger.info("定期重登 @ %d/%d", i, len(todo_snapshots))
|
||
if not _relogin():
|
||
logger.warning("定期重登失败, 继续")
|
||
|
||
try:
|
||
status, rows = download_one_snapshot(index_name, date_str, args.force)
|
||
except Exception as e:
|
||
status, rows = "failed", 0
|
||
logger.debug("download_one_snapshot %s %s 异常: %s", index_name, date_str, e)
|
||
|
||
stats[status] = stats.get(status, 0) + 1
|
||
if status == "ok":
|
||
stats["rows"] += rows
|
||
consec_fail = 0
|
||
elif status == "failed":
|
||
consec_fail += 1
|
||
|
||
# 断路器
|
||
if consec_fail >= CIRCUIT_BREAKER:
|
||
logger.error(
|
||
"[FATAL] 断路器触发: 连续 %d 次失败, baostock 疑似不可达",
|
||
consec_fail,
|
||
)
|
||
circuit_triggered = True
|
||
break
|
||
|
||
# 进度日志
|
||
if (i + 1) % PROGRESS_LOG_EVERY == 0:
|
||
elapsed = time.time() - t_start
|
||
logger.info(
|
||
"进度 %d/%d ok=%d skipped=%d empty=%d failed=%d rows=%d (%.0f秒)",
|
||
i + 1, len(todo_snapshots), stats["ok"], stats["skipped"],
|
||
stats["empty"], stats["failed"], stats["rows"], elapsed,
|
||
)
|
||
|
||
# 限速
|
||
if i < len(todo_snapshots) - 1:
|
||
time.sleep(BS_INTERVAL)
|
||
|
||
# 登出
|
||
try:
|
||
bs.logout()
|
||
except Exception:
|
||
pass
|
||
|
||
elapsed = time.time() - t_start
|
||
logger.info("=" * 60)
|
||
if circuit_triggered:
|
||
logger.info("[RESULT] 断路器触发中止, 耗时 %.1f 秒", elapsed)
|
||
else:
|
||
logger.info("[RESULT] 完成, 耗时 %.1f 秒", elapsed)
|
||
logger.info("统计: %s", json.dumps(stats, ensure_ascii=False))
|
||
|
||
sys.exit(2 if circuit_triggered else 0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|