claude_dev
917d5bca2a
本机下131只后baostock broken pipe卡死(脚本无retry). 加reconnect()+download_one max_retries=3(空结果/broken pipe→bs.logout+login重试). 下次跑断点续传skip 131.
193 lines
7.0 KiB
Python
193 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
||
"""baostock 15min 下载（NAS 容器跑，本机零负载）。
|
||
|
||
为何 baostock：新浪 stock_zh_a_minute 15min 历史仅半年；baostock 5.5 年 + 复权齐全
|
||
（实测浦发 2021-01-04 起 21328 行；除权日 raw -5.9% vs qfq -0.7% 平滑）。
|
||
|
||
双源：qfq(adjustflag=2) + raw(adjustflag=3)。断点续传（skip 已存在，扩范围重下）。
|
||
数据落 NAS 本地盘（容器挂载 /volume1/stock → /stock，最快，不经网络写入）。
|
||
|
||
约束：单线程 + SLEEP 限速（baostock 服务器温和限频，避免封）。
|
||
|
||
用法（NAS 容器，本机编排）：
|
||
/var/packages/Docker/target/usr/bin/docker run --rm -d --name dl15 \\
|
||
--memory=512m -v /volume1/stock:/stock \\
|
||
sanguo_vnpy_v2:with-sqlite \\
|
||
sh -c "pip install baostock -q && python /stock/sanguo_vnpy/scripts/baostock_download.py --all --adjust both --start 2021-01-01"
|
||
|
||
env:
|
||
DL_DIR 输出根（默认 /stock/A股数据/minute_kline）
|
||
STOCK_LIST 全市场 csv（默认 NAS stock_info csv）
|
||
SLEEP 每只每源间隔秒（默认 0.3）
|
||
"""
|
||
import argparse
|
||
import csv
|
||
import logging
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import pandas as pd
|
||
import baostock as bs
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
log = logging.getLogger("dl_15min")
|
||
|
||
DL_DIR = os.environ.get("DL_DIR", "/stock/A股数据/minute_kline")
|
||
SLEEP = float(os.environ.get("SLEEP", "0.3"))
|
||
DEFAULT_STOCK_LIST = "/stock/A股数据/stock_info/stock_basic_info_raw_20260326_113530.csv"
|
||
|
||
|
||
def prefix(code: str) -> str:
|
||
return "sh" if code.startswith(("60", "68", "51", "56", "58")) else "sz"
|
||
|
||
|
||
def bs_symbol(code: str) -> str:
|
||
return f"{prefix(code)}.{code}"
|
||
|
||
|
||
def out_path(code: str, adjust: str) -> str:
|
||
d = os.path.join(DL_DIR, f"15min_{adjust}")
|
||
os.makedirs(d, exist_ok=True)
|
||
return os.path.join(d, f"{prefix(code)}{code}_15min.parquet")
|
||
|
||
|
||
def exists(code: str, adjust: str, start_year: int) -> bool:
|
||
"""断点续传：文件存在且含 start_year 数据（扩范围 start 更早 → 重下覆盖）。"""
|
||
f = out_path(code, adjust)
|
||
if not os.path.exists(f):
|
||
return False
|
||
try:
|
||
df = pd.read_parquet(f, columns=["date"])
|
||
return str(start_year) in df["date"].astype(str).str[:4].unique()
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def reconnect() -> bool:
|
||
"""baostock 断连后重连（限流/broken pipe 恢复）。"""
|
||
try:
|
||
bs.logout()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
return bs.login().error_code == "0"
|
||
|
||
|
||
def download_one(code: str, start: str, end: str, adjustflag: str, max_retries: int = 3):
|
||
"""baostock 15min，返回 (df, None) 或 (None, err)。断连/限流自动 re-login retry。"""
|
||
last_err = ""
|
||
for attempt in range(max_retries):
|
||
try:
|
||
rs = bs.query_history_k_data_plus(
|
||
bs_symbol(code),
|
||
"date,time,open,high,low,close,volume,amount",
|
||
start_date=start, end_date=end, frequency="15", adjustflag=adjustflag,
|
||
)
|
||
rows, fields = [], rs.fields
|
||
while (rs.error_code == "0") & rs.next():
|
||
rows.append(rs.get_row_data())
|
||
if not rows:
|
||
last_err = rs.error_msg or "empty"
|
||
# 空可能是限流（"接收数据异常"）→ re-login 重试
|
||
if attempt < max_retries - 1:
|
||
reconnect()
|
||
time.sleep(2)
|
||
continue
|
||
return None, last_err
|
||
df = pd.DataFrame(rows, columns=fields)
|
||
for c in ("open", "high", "low", "close", "volume", "amount"):
|
||
df[c] = pd.to_numeric(df[c], errors="coerce")
|
||
t = df["time"].astype(str)
|
||
df["datetime"] = pd.to_datetime(
|
||
df["date"] + " " + t.str[8:10] + ":" + t.str[10:12], errors="coerce"
|
||
)
|
||
df = df.dropna(subset=["close", "datetime"]).sort_values("datetime")
|
||
return df, None
|
||
except Exception as e: # noqa: BLE001 broken pipe 等
|
||
last_err = f"{type(e).__name__}: {str(e)[:80]}"
|
||
if attempt < max_retries - 1:
|
||
reconnect()
|
||
time.sleep(2)
|
||
continue
|
||
return None, last_err
|
||
return None, last_err
|
||
|
||
|
||
def save_one(code, df, adjust):
|
||
df.to_parquet(out_path(code, adjust))
|
||
return len(df)
|
||
|
||
|
||
def load_all_codes(stock_list):
|
||
codes = []
|
||
with open(stock_list, encoding="utf-8", errors="replace") as f:
|
||
for row in csv.DictReader(f):
|
||
for k in ("code", "symbol", "ts_code", "代码", "股票代码", "成分券代码"):
|
||
if k in row and row[k]:
|
||
c = str(row[k]).strip().split(".")[0]
|
||
if c.isdigit() and len(c) == 6:
|
||
codes.append(c)
|
||
break
|
||
return codes
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--symbols")
|
||
ap.add_argument("--all", action="store_true")
|
||
ap.add_argument("--start", default="2021-01-01")
|
||
ap.add_argument("--end", default=None)
|
||
ap.add_argument("--adjust", default="qfq", help="qfq / raw / both")
|
||
ap.add_argument("--force", action="store_true")
|
||
args = ap.parse_args()
|
||
|
||
end = args.end or time.strftime("%Y-%m-%d")
|
||
start_year = int(args.start[:4])
|
||
adjusts = ["qfq", "raw"] if args.adjust == "both" else [args.adjust]
|
||
flag = {"qfq": "2", "raw": "3"}
|
||
|
||
if args.symbols:
|
||
codes = [c.strip() for c in args.symbols.split(",") if c.strip()]
|
||
elif args.all:
|
||
sl = os.environ.get("STOCK_LIST", DEFAULT_STOCK_LIST)
|
||
if not os.path.exists(sl):
|
||
log.error("STOCK_LIST 不存在: %s", sl)
|
||
sys.exit(1)
|
||
codes = load_all_codes(sl)
|
||
log.info("--all 读到 %d 只", len(codes))
|
||
else:
|
||
ap.error("需 --symbols 或 --all")
|
||
|
||
lg = bs.login()
|
||
log.info("baostock login: %s %s", lg.error_code, lg.error_msg)
|
||
|
||
ok = fail = rows = skipped = 0
|
||
for i, code in enumerate(codes, 1):
|
||
for adj in adjusts:
|
||
if not args.force and exists(code, adj, start_year):
|
||
skipped += 1
|
||
continue
|
||
df, err = download_one(code, args.start, end, flag[adj])
|
||
if df is None or df.empty:
|
||
fail += 1
|
||
if i <= 20 or i % 100 == 0:
|
||
log.warning("[%d/%d] %s %s empty %s", i, len(codes), code, adj, err)
|
||
else:
|
||
try:
|
||
n = save_one(code, df, adj)
|
||
ok += 1
|
||
rows += n
|
||
if i <= 20 or i % 100 == 0:
|
||
log.info("[%d/%d] %s %s ok %d", i, len(codes), code, adj, n)
|
||
except Exception as e: # noqa: BLE001
|
||
fail += 1
|
||
log.error("save %s %s: %s", code, adj, e)
|
||
time.sleep(SLEEP)
|
||
bs.logout()
|
||
log.info("=== 完成 ok=%d skip=%d fail=%d rows=%d dir=%s ===",
|
||
ok, skipped, fail, rows, DL_DIR)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|