← 返回 下载 dataloader.py
"""数据层:A 股日线行情抓取与本地缓存。

数据源(公开接口,非一手行情,回测用途需自行核对精度):
- 日线(前复权):腾讯行情 web.ifzq.gtimg.cn
- 备用日线:新浪财经 CN_MarketData.getKLineData
- 指数成分股:东方财富 clist 接口(BK0500 = 沪深300)

所有抓取结果按标的缓存到 data/kline/{code}.csv,重复运行不会重复请求。
"""
from __future__ import annotations

import json
import os
import time
import urllib.parse
import urllib.request
from typing import Dict, List, Optional, Tuple

import pandas as pd

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
KLINE_DIR = os.path.join(DATA_DIR, "kline")
UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/120.0 Safari/537.36"


def _http_get(url: str, timeout: int = 20) -> str:
    req = urllib.request.Request(url, headers={"User-Agent": UA, "Referer": "https://gu.qq.com/"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read()
    except Exception:
        # 代理偶发 501/断连时降级为 curl
        import subprocess
        raw = subprocess.run(
            ["curl", "-s", "--max-time", str(timeout), "-A", UA, "-H", "Referer: https://gu.qq.com/", url],
            capture_output=True,
        ).stdout
        if not raw:
            raise
    for enc in ("utf-8", "gbk"):
        try:
            return raw.decode(enc)
        except UnicodeDecodeError:
            continue
    return raw.decode("utf-8", "ignore")


def em_market(code: str) -> str:
    """按 A 股代码规则推断交易所前缀。"""
    code = str(code).zfill(6)
    if code[0] in ("6", "9") or code.startswith("688"):
        return "sh"
    if code[0] in ("0", "3", "2"):
        return "sz"
    if code[0] in ("4", "8"):
        return "bj"
    return "sh"


def fetch_hs300(limit: Optional[int] = None) -> List[Tuple[str, str]]:
    """抓取沪深300成分股列表(快照,非历史成分,存在幸存者偏差)。"""
    url = (
        "https://push2.eastmoney.com/api/qt/clist/get?pn=1&pz=300&po=1&np=1&fltt=2&invt=2"
        "&fid=f3&fs=b:BK0500&fields=f12,f14"
    )
    data = json.loads(_http_get(url))
    diff = data.get("data", {}).get("diff", []) or []
    out = [(d["f12"], d["f14"]) for d in diff if d.get("f12")]
    if limit:
        out = out[:limit]
    return out


def fetch_kline(code: str, start: str = "2021-01-01", end: str = "2030-01-01") -> pd.DataFrame:
    """抓取单标的日线(前复权)。腾讯接口单次约 640 根上限,按 2 年窗口分段抓取合并。"""
    sym = em_market(code) + str(code).zfill(6)
    frames = []
    win_start = pd.Timestamp(start)
    end_ts = min(pd.Timestamp(end), pd.Timestamp.today())
    while win_start < end_ts:
        win_end = min(win_start + pd.DateOffset(years=2) - pd.Timedelta(days=1), end_ts)
        param = f"{sym},day,{win_start.date()},{win_end.date()},2000,qfq"
        url = "https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param=" + urllib.parse.quote(param)
        try:
            payload = json.loads(_http_get(url))
        except Exception:
            # WAF 拦截/网络异常:腾讯路径整体放弃,走新浪兜底
            frames = []
            break
        node = (payload.get("data") or {}).get(sym) or {}
        rows = node.get("qfqday") or node.get("day") or []
        if rows:
            frames.append(pd.DataFrame([r[:6] for r in rows]))
        win_start = win_end + pd.Timedelta(days=1)
    if not frames:
        return _fetch_kline_sina(code)
    df = pd.concat(frames, ignore_index=True).drop_duplicates(subset=0, keep="last")
    df.columns = ["date", "open", "close", "high", "low", "volume"]
    df["date"] = pd.to_datetime(df["date"])
    for col in ["open", "high", "low", "close", "volume"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    df = df.dropna(subset=["close"]).set_index("date").sort_index()
    # 腾讯成交量单位为手,统一换算成股
    df["volume"] = df["volume"] * 100.0
    df.attrs["source"] = "tencent_qfq"
    return df[["open", "high", "low", "close", "volume"]]


def _fetch_kline_sina(code: str) -> pd.DataFrame:
    """新浪兜底源(单次最多约 1600 根,约 6.5 年)。

    注意:新浪该接口为**不复权**价格,除权除息日会出现价格跳空,
    会轻微污染动量/趋势类因子,量化正式研究建议优先使用前复权数据源。
    """
    sym = em_market(code) + str(code).zfill(6)
    surl = (
        "https://money.finance.sina.com.cn/quotes_service/api/json_v2.php/"
        f"CN_MarketData.getKLineData?symbol={sym}&scale=240&ma=no&datalen=1600"
    )
    try:
        rows = json.loads(_http_get(surl))
    except Exception:
        return pd.DataFrame()
    if not rows:
        return pd.DataFrame()
    df = pd.DataFrame(rows)
    df["date"] = pd.to_datetime(df["day"])
    for col in ["open", "high", "low", "close", "volume"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    df = df.dropna(subset=["close"]).set_index("date").sort_index()
    df.attrs["source"] = "sina_unadjusted"
    return df[["open", "high", "low", "close", "volume"]]


def _fetch_sina(code: str, datalen: int = 1200) -> List[list]:
    """新浪备用接口,仅在腾讯失败时使用。"""
    sym = em_market(code) + str(code).zfill(6)
    url = (
        "https://money.finance.sina.com.cn/quotes_service/api/json_v2.php/"
        f"CN_MarketData.getKLineData?symbol={sym}&scale=240&ma=no&datalen={datalen}"
    )
    try:
        rows = json.loads(_http_get(url))
    except Exception:
        return []
    return [
        [r["day"], r["open"], r["close"], r["high"], r["low"], float(r["volume"]) / 100.0]
        for r in rows
    ]


def fetch_index(code: str = "000300", start: str = "2021-01-01", end: str = "2030-01-01") -> pd.DataFrame:
    """抓取指数日线,用于基准对比。腾讯失败时降级新浪。"""
    # 指数前缀规则:399xxx 为深市指数,其余(000300/000905 等)默认沪市
    sym = ("sz" if str(code).startswith("399") else "sh") + str(code).zfill(6)
    param = f"{sym},day,{start},{end},2000,index"
    url = "https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param=" + urllib.parse.quote(param)
    try:
        payload = json.loads(_http_get(url))
        node = (payload.get("data") or {}).get(sym) or {}
        rows = node.get("day") or node.get("qfqday") or []
    except Exception:
        rows = []
    if rows:
        df = pd.DataFrame([r[:6] for r in rows])
        df.columns = ["date", "open", "close", "high", "low", "volume"]
        df["date"] = pd.to_datetime(df["date"])
        for col in ["open", "high", "low", "close", "volume"]:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        return df.dropna(subset=["close"]).set_index("date").sort_index()[["open", "high", "low", "close", "volume"]]
    # 新浪备用(指数同样适用)
    surl = (
        "https://money.finance.sina.com.cn/quotes_service/api/json_v2.php/"
        f"CN_MarketData.getKLineData?symbol={sym}&scale=240&ma=no&datalen=1500"
    )
    rows = json.loads(_http_get(surl))
    if not rows:
        return pd.DataFrame()
    df = pd.DataFrame(rows)
    df["date"] = pd.to_datetime(df["day"])
    for col in ["open", "high", "low", "close"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    return df.dropna(subset=["close"]).set_index("date").sort_index()[["open", "high", "low", "close", "volume"]]


def load_universe(
    codes: List[str],
    start: str,
    end: str,
    use_cache: bool = True,
    verbose: bool = True,
) -> Dict[str, pd.DataFrame]:
    """批量加载日线,带磁盘缓存。"""
    os.makedirs(KLINE_DIR, exist_ok=True)
    # 缓存新鲜度以"今天"为上限判断,否则 end=2030 会导致缓存永不刷新
    target_end = min(pd.Timestamp(end), pd.Timestamp.today().normalize())
    out: Dict[str, pd.DataFrame] = {}
    for i, code in enumerate(codes, 1):
        code = str(code).zfill(6)
        path = os.path.join(KLINE_DIR, f"{code}.csv")
        df = pd.DataFrame()
        if use_cache and os.path.exists(path):
            try:
                cached = pd.read_csv(path, index_col=0, comment="#")
                cached.index = pd.to_datetime(cached.index, errors="coerce")
                cached = cached.dropna()
                if (
                    len(cached)
                    and isinstance(cached.index, pd.DatetimeIndex)
                    and cached.index.max() >= target_end - pd.Timedelta(days=3)
                ):
                    df = cached
            except Exception:
                df = pd.DataFrame()
        if df.empty:
            for attempt in range(3):
                try:
                    df = fetch_kline(code, start, end)
                    if len(df):
                        with open(path, "w", encoding="utf-8") as fh:
                            src = df.attrs.get("source", "")
                            if src:
                                fh.write(f"# source={src}\n")
                            df.to_csv(fh)
                        break
                except Exception as exc:  # 单只失败不影响整体
                    if attempt == 2 and verbose:
                        print(f"  [warn] {code} 抓取失败: {exc}")
                    df = pd.DataFrame()
                time.sleep(0.3)
        if df.empty:
            continue
        if not isinstance(df.index, pd.DatetimeIndex):
            continue
        df = df.loc[(df.index >= pd.Timestamp(start)) & (df.index <= pd.Timestamp(end))]
        if len(df) > 120:
            out[code] = df
        if verbose and i % 25 == 0:
            print(f"  已加载 {i}/{len(codes)}")
    return out