333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""
|
|
Stock data service — wraps AKShare with Redis caching.
|
|
All AKShare calls are blocking I/O; run them in a thread pool via asyncio.
|
|
"""
|
|
import asyncio
|
|
import json
|
|
from functools import lru_cache
|
|
from typing import Optional
|
|
from loguru import logger
|
|
|
|
try:
|
|
import akshare as ak
|
|
AK_AVAILABLE = True
|
|
except Exception:
|
|
AK_AVAILABLE = False
|
|
logger.warning("AKShare not available, using mock data")
|
|
|
|
|
|
# ── helpers ──────────────────────────────────────────────────────────────────
|
|
|
|
def _run_sync(fn, *args, **kwargs):
|
|
"""Run a sync blocking function in the default thread pool."""
|
|
loop = asyncio.get_running_loop()
|
|
return loop.run_in_executor(None, lambda: fn(*args, **kwargs))
|
|
|
|
|
|
# ── market overview ───────────────────────────────────────────────────────────
|
|
|
|
MAJOR_INDICES = [
|
|
("sh000001", "上证指数"),
|
|
("sz399001", "深证成指"),
|
|
("sz399006", "创业板指"),
|
|
("sh000688", "科创50"),
|
|
("sh000300", "沪深300"),
|
|
]
|
|
|
|
|
|
async def get_market_overview() -> list[dict]:
|
|
if not AK_AVAILABLE:
|
|
return _mock_market_overview()
|
|
|
|
try:
|
|
df = await _run_sync(ak.stock_zh_index_spot_em)
|
|
result = []
|
|
code_map = {code: name for code, name in MAJOR_INDICES}
|
|
for _, row in df.iterrows():
|
|
code = str(row.get("代码", ""))
|
|
if code in code_map:
|
|
result.append({
|
|
"index_code": code,
|
|
"index_name": code_map[code],
|
|
"current": float(row.get("最新价", 0)),
|
|
"change": float(row.get("涨跌额", 0)),
|
|
"change_pct": float(row.get("涨跌幅", 0)),
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"get_market_overview error: {e}")
|
|
return _mock_market_overview()
|
|
|
|
|
|
# ── real-time quotes (A-share spot) ──────────────────────────────────────────
|
|
|
|
async def get_all_stocks_spot() -> list[dict]:
|
|
"""All A-share real-time quotes — used for heatmap."""
|
|
if not AK_AVAILABLE:
|
|
return _mock_heatmap_data()
|
|
|
|
try:
|
|
df = await _run_sync(ak.stock_zh_a_spot_em)
|
|
result = []
|
|
for _, row in df.iterrows():
|
|
pct = float(row.get("涨跌幅", 0) or 0)
|
|
result.append({
|
|
"symbol": str(row.get("代码", "")),
|
|
"name": str(row.get("名称", "")),
|
|
"price": float(row.get("最新价", 0) or 0),
|
|
"change": float(row.get("涨跌额", 0) or 0),
|
|
"change_pct": pct,
|
|
"open": float(row.get("今开", 0) or 0),
|
|
"high": float(row.get("最高", 0) or 0),
|
|
"low": float(row.get("最低", 0) or 0),
|
|
"prev_close": float(row.get("昨收", 0) or 0),
|
|
"volume": float(row.get("成交量", 0) or 0),
|
|
"amount": float(row.get("成交额", 0) or 0),
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"get_all_stocks_spot error: {e}")
|
|
return _mock_heatmap_data()
|
|
|
|
|
|
async def get_stock_quote(symbol: str) -> Optional[dict]:
|
|
"""Single stock real-time quote."""
|
|
all_stocks = await get_all_stocks_spot()
|
|
for s in all_stocks:
|
|
if s["symbol"] == symbol:
|
|
return s
|
|
return None
|
|
|
|
|
|
# ── K-line data ───────────────────────────────────────────────────────────────
|
|
|
|
PERIOD_MAP = {
|
|
"daily": "daily",
|
|
"weekly": "weekly",
|
|
"monthly": "monthly",
|
|
}
|
|
|
|
|
|
async def get_kline(symbol: str, period: str = "daily", adjust: str = "qfq", limit: int = 250) -> list[dict]:
|
|
"""
|
|
period: daily | weekly | monthly
|
|
adjust: qfq (前复权) | hfq (后复权) | "" (不复权)
|
|
"""
|
|
if not AK_AVAILABLE:
|
|
return _mock_kline(limit)
|
|
|
|
ak_period = PERIOD_MAP.get(period, "daily")
|
|
|
|
try:
|
|
df = await _run_sync(
|
|
ak.stock_zh_a_hist,
|
|
symbol=symbol,
|
|
period=ak_period,
|
|
adjust=adjust,
|
|
)
|
|
df = df.tail(limit)
|
|
result = []
|
|
for _, row in df.iterrows():
|
|
result.append({
|
|
"date": str(row.get("日期", "")),
|
|
"open": float(row.get("开盘", 0)),
|
|
"high": float(row.get("最高", 0)),
|
|
"low": float(row.get("最低", 0)),
|
|
"close": float(row.get("收盘", 0)),
|
|
"volume": float(row.get("成交量", 0)),
|
|
"amount": float(row.get("成交额", 0)),
|
|
"change_pct": float(row.get("涨跌幅", 0)),
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"get_kline {symbol} {period} error: {e}")
|
|
return _mock_kline(limit)
|
|
|
|
|
|
async def get_intraday(symbol: str) -> list[dict]:
|
|
"""Today's minute-level data."""
|
|
if not AK_AVAILABLE:
|
|
return _mock_intraday()
|
|
|
|
try:
|
|
df = await _run_sync(ak.stock_zh_a_hist_min_em, symbol=symbol, period="1", adjust="")
|
|
result = []
|
|
for _, row in df.iterrows():
|
|
result.append({
|
|
"time": str(row.get("时间", "")),
|
|
"price": float(row.get("收盘", 0)),
|
|
"volume": float(row.get("成交量", 0)),
|
|
"amount": float(row.get("成交额", 0)),
|
|
"avg_price": float(row.get("均价", 0) or 0),
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"get_intraday {symbol} error: {e}")
|
|
return _mock_intraday()
|
|
|
|
|
|
async def get_five_day(symbol: str) -> list[dict]:
|
|
"""5-day minute-level data."""
|
|
if not AK_AVAILABLE:
|
|
return _mock_intraday(days=5)
|
|
|
|
try:
|
|
df = await _run_sync(ak.stock_zh_a_hist_min_em, symbol=symbol, period="1", adjust="")
|
|
result = []
|
|
for _, row in df.iterrows():
|
|
result.append({
|
|
"time": str(row.get("时间", "")),
|
|
"price": float(row.get("收盘", 0)),
|
|
"volume": float(row.get("成交量", 0)),
|
|
"amount": float(row.get("成交额", 0) or 0),
|
|
"avg_price": float(row.get("均价", 0) or 0),
|
|
})
|
|
return result[-5 * 240:]
|
|
except Exception as e:
|
|
logger.error(f"get_five_day {symbol} error: {e}")
|
|
return _mock_intraday(days=5)
|
|
|
|
|
|
# ── search ────────────────────────────────────────────────────────────────────
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _get_stock_list_cached() -> list[dict]:
|
|
"""Cache the full stock list in memory (refreshed on process restart)."""
|
|
if not AK_AVAILABLE:
|
|
return []
|
|
try:
|
|
import akshare as ak
|
|
df = ak.stock_info_a_code_name()
|
|
return [{"symbol": str(r["code"]), "name": str(r["name"]), "market": "A股"} for _, r in df.iterrows()]
|
|
except Exception as e:
|
|
logger.error(f"_get_stock_list_cached error: {e}")
|
|
return []
|
|
|
|
|
|
async def search_stocks(query: str, limit: int = 20) -> list[dict]:
|
|
stock_list = await _run_sync(_get_stock_list_cached)
|
|
query = query.strip().lower()
|
|
results = [
|
|
s for s in stock_list
|
|
if query in s["symbol"].lower() or query in s["name"].lower()
|
|
]
|
|
return results[:limit]
|
|
|
|
|
|
# ── sector data for heatmap ───────────────────────────────────────────────────
|
|
|
|
async def get_sector_spot() -> list[dict]:
|
|
"""Board/sector change pct for heatmap grouping."""
|
|
if not AK_AVAILABLE:
|
|
return _mock_sectors()
|
|
|
|
try:
|
|
df = await _run_sync(ak.stock_board_industry_name_em)
|
|
result = []
|
|
for _, row in df.iterrows():
|
|
result.append({
|
|
"sector": str(row.get("板块名称", "")),
|
|
"change_pct": float(row.get("涨跌幅", 0) or 0),
|
|
"volume": float(row.get("成交量", 0) or 0),
|
|
"amount": float(row.get("成交额", 0) or 0),
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"get_sector_spot error: {e}")
|
|
return _mock_sectors()
|
|
|
|
|
|
# ── mock data (fallback when market is closed or AKShare unavailable) ─────────
|
|
|
|
import random
|
|
import math
|
|
from datetime import date, timedelta
|
|
|
|
|
|
def _mock_market_overview() -> list[dict]:
|
|
return [
|
|
{"index_code": "sh000001", "index_name": "上证指数", "current": 3312.46, "change": -28.84, "change_pct": -0.86},
|
|
{"index_code": "sz399001", "index_name": "深证成指", "current": 10573.99, "change": -93.16, "change_pct": -0.87},
|
|
{"index_code": "sz399006", "index_name": "创业板指", "current": 2105.37, "change": -18.42, "change_pct": -0.87},
|
|
{"index_code": "sh000688", "index_name": "科创50", "current": 968.12, "change": -9.56, "change_pct": -0.98},
|
|
{"index_code": "sh000300", "index_name": "沪深300", "current": 3843.20, "change": -30.11, "change_pct": -0.78},
|
|
]
|
|
|
|
|
|
def _mock_heatmap_data() -> list[dict]:
|
|
sectors = ["银行", "电力设备", "食品饮料", "医药生物", "电子", "汽车", "非银金融", "计算机", "有色金属", "化工"]
|
|
stocks = []
|
|
for i in range(80):
|
|
pct = round(random.uniform(-5, 5), 2)
|
|
sector = sectors[i % len(sectors)]
|
|
stocks.append({
|
|
"symbol": f"{600000 + i:06d}",
|
|
"name": f"测试股票{i+1:02d}",
|
|
"price": round(random.uniform(5, 100), 2),
|
|
"change": round(pct * 0.1, 2),
|
|
"change_pct": pct,
|
|
"open": round(random.uniform(5, 100), 2),
|
|
"high": round(random.uniform(5, 100), 2),
|
|
"low": round(random.uniform(5, 100), 2),
|
|
"prev_close": round(random.uniform(5, 100), 2),
|
|
"volume": random.randint(100000, 10000000),
|
|
"amount": random.randint(1000000, 100000000),
|
|
"sector": sector,
|
|
})
|
|
return stocks
|
|
|
|
|
|
def _mock_kline(limit: int = 250) -> list[dict]:
|
|
bars = []
|
|
price = 20.0
|
|
today = date.today()
|
|
for i in range(limit):
|
|
d = today - timedelta(days=limit - i)
|
|
pct = random.uniform(-0.05, 0.05)
|
|
close = round(price * (1 + pct), 2)
|
|
high = round(max(price, close) * random.uniform(1.0, 1.03), 2)
|
|
low = round(min(price, close) * random.uniform(0.97, 1.0), 2)
|
|
bars.append({
|
|
"date": d.isoformat(),
|
|
"open": round(price, 2),
|
|
"high": high,
|
|
"low": low,
|
|
"close": close,
|
|
"volume": random.randint(500000, 5000000),
|
|
"amount": random.randint(5000000, 50000000),
|
|
"change_pct": round(pct * 100, 2),
|
|
})
|
|
price = close
|
|
return bars
|
|
|
|
|
|
def _mock_intraday(days: int = 1) -> list[dict]:
|
|
bars = []
|
|
price = 20.0
|
|
for d in range(days):
|
|
for minute in range(240):
|
|
h = 9 + minute // 60
|
|
m = (minute % 60) + (30 if d == 0 and minute < 60 else 0)
|
|
if h == 11 and m >= 30:
|
|
continue
|
|
pct = random.uniform(-0.01, 0.01)
|
|
price = round(price * (1 + pct), 2)
|
|
bars.append({
|
|
"time": f"2026-06-0{d+1} {h:02d}:{m % 60:02d}",
|
|
"price": price,
|
|
"volume": random.randint(10000, 500000),
|
|
"amount": random.randint(100000, 5000000),
|
|
"avg_price": price,
|
|
})
|
|
return bars
|
|
|
|
|
|
def _mock_sectors() -> list[dict]:
|
|
sectors = [
|
|
"银行", "电力设备", "食品饮料", "医药生物", "电子",
|
|
"汽车", "非银金融", "计算机", "有色金属", "化工",
|
|
"机械设备", "建筑材料", "传媒", "房地产", "交通运输",
|
|
]
|
|
return [{"sector": s, "change_pct": round(random.uniform(-3, 3), 2), "volume": 0, "amount": 0}
|
|
for s in sectors]
|