146 lines
5.6 KiB
Python
146 lines
5.6 KiB
Python
"""路由包共享件：JSON 直返缓存、复权换算、行转 Bar、共享常量与 SQL。
|
||
|
||
各域路由模块（stocks/etfs/market/backtest/screener/user）从这里取公共工具，
|
||
域内私有工具留在各自文件里。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import bisect
|
||
|
||
import pandas as pd
|
||
from fastapi import Response
|
||
from sqlalchemy import text
|
||
|
||
from .. import cache
|
||
from ..domain import Bar
|
||
|
||
# 复权模式白名单
|
||
ADJUST_MODES = ("bfq", "qfq", "hfq")
|
||
# MA 全量集合（前端已改为本地计算 MA，后端始终返回此集合以保证缓存一致）
|
||
FULL_MA_SET = (5, 10, 20, 30, 60, 120, 250)
|
||
# 指数 K 线支持的周期（日线基底聚合）
|
||
INDEX_TIMEFRAMES = ("1d", "1w", "1M", "1y")
|
||
|
||
|
||
def raw_json(resp) -> str:
|
||
"""pydantic-core（Rust）序列化：与 response_model 直返时的字节完全一致（紧凑分隔符、
|
||
非 ASCII 直出、浮点小数形式），且比 stdlib json.dumps 快。大响应（preview ~250KB）
|
||
命中缓存时直接 Response 原样返回，跳过校验/再序列化。"""
|
||
return resp.model_dump_json()
|
||
|
||
|
||
async def cached_json_response(key: str) -> Response | None:
|
||
"""两级缓存读（进程内 → Redis）：命中返回可直接吐给客户端的 Response。
|
||
存的均为序列化好的 JSON 字符串（Redis 侧 json.loads 后仍是 str），Redis 命中顺手晋级本地。"""
|
||
raw = cache.local_get(key)
|
||
if raw is None:
|
||
raw = await cache.cache_get(key)
|
||
if not isinstance(raw, str):
|
||
return None
|
||
cache.local_set(key, raw, ttl=120)
|
||
return Response(content=raw, media_type="application/json")
|
||
|
||
|
||
def series_to_jsonable(s: pd.Series) -> list[float | None]:
|
||
"""NaN -> None（lightweight-charts 的 whitespace data，跳过指标预热期）。"""
|
||
out: list[float | None] = []
|
||
for v in s.tolist():
|
||
if v is None or (isinstance(v, float) and v != v):
|
||
out.append(None)
|
||
else:
|
||
out.append(float(v))
|
||
return out
|
||
|
||
|
||
def rows_to_bars(rows) -> list[Bar]:
|
||
return [
|
||
Bar(
|
||
ts=r.ts, open=r.open, high=r.high, low=r.low, close=r.close, volume=r.volume,
|
||
amount=getattr(r, "amount", None), turnover=getattr(r, "turnover", None),
|
||
)
|
||
for r in rows
|
||
]
|
||
|
||
|
||
# 信息卡一条 SQL 拿全：stock_basic 基本信息 + 「优先与行情同日、缺则最新日」的 daily_snapshot
|
||
# （LATERAL 单条替换原两条查询，语义不变：target 为 NULL 时全按最新日兜底）。
|
||
# ETF 走 etf_basic 分支（代码前缀与股票不重叠，两分支至多一个命中）：
|
||
# 名称/上市日来自表内，市值（元）换算成万元与快照口径一致，无 PE/PB。
|
||
INFO_SQL = text(
|
||
"""
|
||
SELECT ts_code, symbol, name, industry, area, market, list_date,
|
||
turnover_rate, pe_ttm, pb, total_mv, circ_mv
|
||
FROM (
|
||
SELECT sb.ts_code, sb.symbol, sb.name, sb.industry, sb.area, sb.market, sb.list_date,
|
||
ds.turnover_rate, ds.pe_ttm, ds.pb, ds.total_mv, ds.circ_mv
|
||
FROM stock_basic sb
|
||
LEFT JOIN LATERAL (
|
||
SELECT turnover_rate, pe_ttm, pb, total_mv, circ_mv
|
||
FROM daily_snapshot
|
||
WHERE ts_code = sb.ts_code
|
||
ORDER BY (trade_date = cast(:target AS timestamp)) DESC, trade_date DESC
|
||
LIMIT 1
|
||
) ds ON true
|
||
WHERE sb.ts_code = :code
|
||
UNION ALL
|
||
SELECT eb.ts_code, eb.symbol, eb.name, NULL, NULL,
|
||
CASE eb.exchange WHEN 'SH' THEN '沪市' ELSE '深市' END, eb.list_date,
|
||
eb.turnover_rate, NULL, NULL,
|
||
eb.total_mv / 10000.0, eb.circ_mv / 10000.0
|
||
FROM etf_basic eb
|
||
WHERE eb.ts_code = :code
|
||
) t
|
||
LIMIT 1
|
||
"""
|
||
)
|
||
|
||
# 复权因子是阶梯函数（除权日之间不变）：只取「变化点」行，把每符号 ~7000 行日级因子压到
|
||
# 几十行（600118 仅 32 行），传输量再降两个数量级；lag 窗口在覆盖索引上走 Index Only Scan。
|
||
# upto 传全局最新交易日（非分页）或 end 日期（分页）；窗口首行 prev 为 NULL 恒被保留（窗口基线因子）。
|
||
FACTOR_STEP_SQL = text(
|
||
"""
|
||
SELECT trade_date, adj_factor FROM (
|
||
SELECT trade_date, adj_factor,
|
||
lag(adj_factor) OVER (ORDER BY trade_date) AS prev
|
||
FROM adj_factor
|
||
WHERE ts_code = :code AND trade_date <= :upto
|
||
) t
|
||
WHERE adj_factor IS DISTINCT FROM prev
|
||
ORDER BY trade_date
|
||
"""
|
||
)
|
||
|
||
|
||
def adjust_bars(bars: list[Bar], factors, from_mode: str, to_mode: str) -> list[Bar]:
|
||
"""按复权因子把 K 线从 from_mode 换算到 to_mode（bfq/qfq/hfq）。
|
||
|
||
相对不复权的乘数：bfq=1，qfq=f(t)/f(latest)，hfq=f(t)。
|
||
因子缺失的日期向前沿用最近因子（因子是阶梯函数，除权日之间不变）。
|
||
"""
|
||
fd = sorted((f[0].date(), float(f[1])) for f in factors)
|
||
fdates = [d for d, _ in fd]
|
||
f_latest = fd[-1][1]
|
||
|
||
def _f_at(d) -> float:
|
||
i = bisect.bisect_right(fdates, d) - 1
|
||
return fd[i][1] if i >= 0 else fd[0][1]
|
||
|
||
def _mult(mode: str, f: float) -> float:
|
||
if mode == "bfq":
|
||
return 1.0
|
||
return f / f_latest if mode == "qfq" else f
|
||
|
||
out: list[Bar] = []
|
||
for b in bars:
|
||
f = _f_at(b.ts.date())
|
||
m = _mult(to_mode, f) / _mult(from_mode, f)
|
||
out.append(Bar(
|
||
ts=b.ts,
|
||
open=round(b.open * m, 3), high=round(b.high * m, 3),
|
||
low=round(b.low * m, 3), close=round(b.close * m, 3),
|
||
volume=b.volume,
|
||
# 成交额/换手率是名义量，不随复权缩放
|
||
amount=b.amount, turnover=b.turnover,
|
||
))
|
||
return out
|