129 lines
5.1 KiB
Python
129 lines
5.1 KiB
Python
"""每日热点数据采集：核心股前100 + 题材涨幅前10，收盘后自动入库
|
||
|
||
由 main.lifespan 启动后台任务；幂等（按交易日 UNIQUE 去重）。
|
||
"""
|
||
|
||
import asyncio
|
||
import traceback
|
||
from datetime import datetime, time as dtime, timezone, timedelta
|
||
from typing import Optional
|
||
|
||
from database import get_connection
|
||
from services.themes import fetch_theme_list
|
||
|
||
_CST = timezone(timedelta(hours=8))
|
||
|
||
# 每天采集的后台任务：每 300 秒（5 分钟）检查一次
|
||
CHECK_INTERVAL_SECONDS = 300
|
||
COLLECT_AFTER_TIME = dtime(15, 0) # 收盘后 15:00 开始允许采集
|
||
CORE_STOCK_LIMIT = 100 # 核心股前100
|
||
TOP_THEME_LIMIT = 10 # 题材前10
|
||
|
||
|
||
def _is_trading_day(d: datetime) -> bool:
|
||
"""仅按工作日判断：周一至周五视为交易日，不处理法定节假日"""
|
||
return d.weekday() < 5
|
||
|
||
|
||
def _has_collected(trade_date: str) -> bool:
|
||
"""当日核心股是否已采集"""
|
||
conn = get_connection()
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT 1 FROM daily_core_stocks WHERE trade_date = ? LIMIT 1",
|
||
(trade_date,),
|
||
).fetchone()
|
||
return row is not None
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
async def collect_daily(trade_date: str, dry_run: bool = False) -> dict:
|
||
"""采集指定交易日数据并入库。
|
||
|
||
Args:
|
||
trade_date: YYYY-MM-DD
|
||
dry_run: True 只打印不写库（用于验证）
|
||
|
||
Returns:
|
||
{"core_count": int, "theme_count": int, "skipped": bool}
|
||
"""
|
||
if _has_collected(trade_date):
|
||
print(f"[collector] {trade_date} 已采集，跳过")
|
||
return {"core_count": 0, "theme_count": 0, "skipped": True}
|
||
|
||
# 1. 拉取全部题材列表（含领涨股，作为当日全部股票的采样来源）
|
||
themes = await fetch_theme_list(1, False)
|
||
if not themes:
|
||
print(f"[collector] {trade_date} 题材列表为空（东财失败），跳过")
|
||
return {"core_count": 0, "theme_count": 0, "skipped": True}
|
||
|
||
# 2. 核心股前100：按领涨股 f3 降序，去重（同一股票可能是多个题材领涨股）
|
||
stock_map: dict[str, dict] = {}
|
||
for t in themes:
|
||
code = t.get("securityCode")
|
||
if not code:
|
||
continue
|
||
if code not in stock_map or (t.get("f3") or 0) > (stock_map[code].get("f3") or 0):
|
||
stock_map[code] = {
|
||
"stock_code": code,
|
||
"stock_name": t.get("securityName", ""),
|
||
"f3": t.get("f3"),
|
||
}
|
||
core_stocks = sorted(
|
||
stock_map.values(), key=lambda x: -(x["f3"] or 0)
|
||
)[:CORE_STOCK_LIMIT]
|
||
|
||
# 3. 题材前10：bf3 降序
|
||
top_themes = sorted(themes, key=lambda x: -(x.get("bf3") or 0))[:TOP_THEME_LIMIT]
|
||
|
||
if dry_run:
|
||
print(f"[collector] {trade_date} 核心股 {len(core_stocks)} 只，题材前10 {len(top_themes)} 只")
|
||
return {"core_count": len(core_stocks), "theme_count": len(top_themes), "skipped": False}
|
||
|
||
# 4. 入库（事务，UNIQUE 幂等）
|
||
conn = get_connection()
|
||
try:
|
||
for i, s in enumerate(core_stocks, start=1):
|
||
conn.execute(
|
||
"INSERT OR IGNORE INTO daily_core_stocks (trade_date, stock_code, stock_name, f3, rank) VALUES (?,?,?,?,?)",
|
||
(trade_date, s["stock_code"], s["stock_name"], s["f3"], i),
|
||
)
|
||
# 核心股所属题材：从 themes 列表（含 securityCode/themeCode/themeName）中
|
||
# 为每个核心股收集其全部所属题材，写入 daily_core_stock_themes
|
||
core_codes = {s["stock_code"] for s in core_stocks}
|
||
for t in themes:
|
||
if t.get("securityCode") in core_codes:
|
||
conn.execute(
|
||
"INSERT OR IGNORE INTO daily_core_stock_themes (trade_date, stock_code, theme_code, theme_name) VALUES (?,?,?,?)",
|
||
(trade_date, t["securityCode"], t["themeCode"], t["themeName"]),
|
||
)
|
||
for i, t in enumerate(top_themes, start=1):
|
||
conn.execute(
|
||
"INSERT OR IGNORE INTO daily_top_themes (trade_date, theme_code, theme_name, bf3, hot_rank, rank) VALUES (?,?,?,?,?,?)",
|
||
(trade_date, t["themeCode"], t["themeName"], t.get("bf3"), t.get("hotRank"), i),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
print(f"[collector] {trade_date} 已采集：核心股 {len(core_stocks)} 只，题材前10 {len(top_themes)} 只")
|
||
return {"core_count": len(core_stocks), "theme_count": len(top_themes), "skipped": False}
|
||
|
||
|
||
async def collector_loop(stop: Optional[asyncio.Event] = None) -> None:
|
||
"""后台循环：每个交易日 15:00 后自动采集当日数据（幂等）"""
|
||
while True:
|
||
try:
|
||
now = datetime.now(_CST)
|
||
if _is_trading_day(now) and now.time() >= COLLECT_AFTER_TIME:
|
||
trade_date = now.strftime("%Y-%m-%d")
|
||
if not _has_collected(trade_date):
|
||
await collect_daily(trade_date)
|
||
except Exception:
|
||
print("[collector] 采集异常：")
|
||
traceback.print_exc()
|
||
if stop is not None and stop.is_set():
|
||
break
|
||
await asyncio.sleep(CHECK_INTERVAL_SECONDS)
|