FinDataHub/Scripts/backfill_akshare.py

#!/usr/bin/env python3
"""补全 akshare stock_daily 历史数据（缺口回补）。
从 tushare_data.stock_basic 取全市场代码（绕开 akshare 全市场列表返回不全的问题），
一次性拉取每只股票的完整日期范围，比逐日拉取快 10 倍。
"""
import sys, time, os, math
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

import akshare as ak
import pandas as pd
import pymysql

DB_CFG = {
    "host": "127.0.0.1", "port": 3306, "user": "root",
    "password": os.environ.get("MYSQL_PASSWORD", ""), "database": "akshare_data", "charset": "utf8mb4",
}
START_DATE = "20260811"
END_DATE = "20260826"
BATCH_SIZE = 50
MAX_WORKERS = 1  # libmini_racer 不是线程安全的，必须单线程

def load_codes_from_db():
    """从 tushare_data.stock_basic 读取全市场 (code, exchange) 列表."""
    db = dict(DB_CFG); db["database"] = "tushare_data"
    conn = pymysql.connect(**db)
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT ts_code FROM stock_basic WHERE ts_code LIKE '%.%'")
            out = []
            for (ts_code,) in cur.fetchall():
                code, exch = ts_code.split(".")
                if exch in ("SH", "SZ"):
                    out.append((code, exch))
            return out
    finally:
        conn.close()

def fetch_stock(code, exch):
    """拉取单只股票的日线数据"""
    try:
        pre = "sh" if exch == "SH" else "sz"
        d = ak.stock_zh_a_daily(symbol=f"{pre}{code}", start_date=START_DATE, end_date=END_DATE)
        if d is None or d.empty:
            return None
        d.insert(0, "ts_code", f"{code}.{exch}")
        d.rename(columns={"date": "trade_date", "volume": "vol"}, inplace=True)
        d["pre_close"] = d["close"].shift(1)
        d["pct_chg"] = round((d["close"] - d["pre_close"]) / d["pre_close"] * 100, 2)
        # 将 NaN 替换为 None（pymysql 不支持 NaN）
        d = d.where(pd.notna(d), None)
        cols = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "vol", "amount", "pct_chg"]
        return d[cols]
    except Exception:
        return None

def upsert_batch(conn, df):
    """批量 upsert 写入"""
    if df is None or df.empty:
        return 0
    cols = list(df.columns)
    cols_str = ", ".join(f"`{c}`" for c in cols)
    placeholders = ", ".join(["%s"] * len(cols))
    update_str = ", ".join(f"`{c}`=VALUES(`{c}`)" for c in cols if c not in ("ts_code", "trade_date"))
    sql = f"INSERT INTO stock_daily ({cols_str}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {update_str}"
    with conn.cursor() as cur:
        for _, row in df.iterrows():
            vals = tuple(row[c] if not (isinstance(row[c], float) and math.isnan(row[c])) else None for c in cols)
            cur.execute(sql, vals)
    conn.commit()
    return len(df)

def main():
    t0 = time.time()
    print(f"补数据范围: {START_DATE} ~ {END_DATE}", flush=True)

    # 获取全市场股票列表（来自 tushare_data.stock_basic，避免 akshare 列表返回不全）
    codes = load_codes_from_db()
    print(f"总股票数: {len(codes)}", flush=True)

    conn = pymysql.connect(**DB_CFG)
    ok, fail, total_rows = 0, 0, 0
    batch = []

    # 单线程顺序执行（libmini_racer 不支持多线程）
    for i, (code, exch) in enumerate(codes, 1):
        result = fetch_stock(code, exch)
        if result is not None:
            batch.append(result)
            ok += 1
        else:
            fail += 1

        if len(batch) >= BATCH_SIZE:
            combined = pd.concat(batch, ignore_index=True)
            written = upsert_batch(conn, combined)
            total_rows += written
            batch.clear()
            print(f"  [{i}/{len(codes)}] ok={ok} fail={fail} rows={total_rows} ({(time.time()-t0)/60:.1f}分)", flush=True)
        elif i % 100 == 0:
            print(f"  [{i}/{len(codes)}] ok={ok} fail={fail} rows={total_rows} ({(time.time()-t0)/60:.1f}分)", flush=True)

    if batch:
        combined = pd.concat(batch, ignore_index=True)
        written = upsert_batch(conn, combined)
        total_rows += written

    conn.close()
    elapsed = (time.time() - t0) / 60
    print(f"完成: {ok}成功/{fail}失败, {total_rows}行写入, 耗时{elapsed:.1f}分", flush=True)

if __name__ == "__main__":
    main()