97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
"""A 股指数全量日线（上证指数日 K 图数据源）。
|
||
|
||
candles 表只有 6 位纯代码股票（TDX 导入明确排除指数，sh000001 与 sz000001 平安银行
|
||
无法区分），指数走 tushare index_daily 按需拉取：分页拉全量（1990 年至今 ~8900 根），
|
||
进程内 + Redis 两级缓存，历史不可变、TTL 兜到当日更新。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import time
|
||
from datetime import datetime
|
||
|
||
from .. import cache
|
||
from ..domain import Bar
|
||
|
||
_PAGE = 8000 # index_daily 单次返回上限
|
||
_TTL = 7200 # 缓存 2h：历史不可变，只影响当日 bar 的新鲜度
|
||
_CALL_GAP = 0.12 # 分页请求间隔（秒），对 tushare 控频（与 market_overview 同款）
|
||
|
||
SH_INDEX = "000001.SH"
|
||
|
||
# 进程内缓存：ts_code -> (bars, 过期时刻)。bars 为全量日线（升序）
|
||
_mem: dict[str, tuple[list[Bar], float]] = {}
|
||
|
||
|
||
def _fetch_all_sync(ts_code: str) -> list[Bar]:
|
||
"""分页拉全量日线。vol 单位手 -> 股，amount 千元 -> 元（与 fetch_daily 同款换算）。"""
|
||
from .tushare_provider import get_pro
|
||
|
||
pro = get_pro()
|
||
frames = []
|
||
offset = 0
|
||
while True:
|
||
df = pro.index_daily(ts_code=ts_code, offset=offset, limit=_PAGE)
|
||
if df is None or df.empty:
|
||
break
|
||
frames.append(df)
|
||
if len(df) < _PAGE:
|
||
break
|
||
offset += _PAGE
|
||
time.sleep(_CALL_GAP)
|
||
if not frames:
|
||
raise RuntimeError(f"Tushare index_daily 无数据: {ts_code}")
|
||
|
||
import pandas as pd
|
||
|
||
df = pd.concat(frames).drop_duplicates(subset="trade_date").sort_values("trade_date")
|
||
bars: list[Bar] = []
|
||
for _, r in df.iterrows():
|
||
amt = r.get("amount")
|
||
bars.append(
|
||
Bar(
|
||
ts=datetime.strptime(str(r["trade_date"]), "%Y%m%d"),
|
||
open=float(r["open"]), high=float(r["high"]),
|
||
low=float(r["low"]), close=float(r["close"]),
|
||
volume=float(r["vol"]) * 100.0,
|
||
amount=float(amt) * 1000.0 if amt is not None and amt == amt else None,
|
||
)
|
||
)
|
||
return bars
|
||
|
||
|
||
def _to_raw(bars: list[Bar]) -> str:
|
||
"""紧凑 JSON：[[ts_iso, open, high, low, close, volume, amount|null], ...]（~600KB）。"""
|
||
return json.dumps(
|
||
[[b.ts.isoformat(), b.open, b.high, b.low, b.close, b.volume, b.amount] for b in bars],
|
||
ensure_ascii=False, separators=(",", ":"),
|
||
)
|
||
|
||
|
||
def _from_raw(raw: str) -> list[Bar]:
|
||
return [
|
||
Bar(ts=datetime.fromisoformat(row[0]), open=row[1], high=row[2], low=row[3],
|
||
close=row[4], volume=row[5], amount=row[6])
|
||
for row in json.loads(raw)
|
||
]
|
||
|
||
|
||
async def get_index_daily(ts_code: str = SH_INDEX) -> list[Bar]:
|
||
"""全量日线（升序）：进程内 -> Redis -> tushare，未命中层级回填上一级。"""
|
||
hit = _mem.get(ts_code)
|
||
if hit and hit[1] > time.monotonic():
|
||
return hit[0]
|
||
|
||
key = f"idxd:{ts_code}"
|
||
raw = await cache.cache_get(key)
|
||
if isinstance(raw, str):
|
||
bars = _from_raw(raw)
|
||
_mem[ts_code] = (bars, time.monotonic() + _TTL)
|
||
return bars
|
||
|
||
bars = await asyncio.to_thread(_fetch_all_sync, ts_code)
|
||
_mem[ts_code] = (bars, time.monotonic() + _TTL)
|
||
await cache.cache_set(key, _to_raw(bars), ttl=_TTL)
|
||
return bars
|