223 lines
7.9 KiB
Python
223 lines
7.9 KiB
Python
"""自选股 API。"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
import time
|
|
from datetime import date
|
|
|
|
import polars as pl
|
|
from fastapi import APIRouter, Query, Request
|
|
from pydantic import BaseModel
|
|
|
|
from app.services import watchlist
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/watchlist", tags=["watchlist"])
|
|
|
|
|
|
class AddRequest(BaseModel):
|
|
symbol: str
|
|
note: str = ""
|
|
|
|
|
|
class BatchAddRequest(BaseModel):
|
|
symbols: list[str]
|
|
note: str = ""
|
|
|
|
|
|
def _with_names(rows: list[dict], request: Request) -> list[dict]:
|
|
if not rows:
|
|
return rows
|
|
try:
|
|
df_i = request.app.state.repo.get_instruments()
|
|
if df_i.is_empty() or "symbol" not in df_i.columns or "name" not in df_i.columns:
|
|
return rows
|
|
name_by_symbol = dict(df_i.select(["symbol", "name"]).iter_rows())
|
|
return [{**row, "name": name_by_symbol.get(row.get("symbol"))} for row in rows]
|
|
except Exception as e: # noqa: BLE001
|
|
logger.debug("attach watchlist names failed: %s", e)
|
|
return rows
|
|
|
|
|
|
@router.get("")
|
|
def list_all(request: Request):
|
|
return {"symbols": _with_names(watchlist.list_symbols(), request)}
|
|
|
|
|
|
@router.post("")
|
|
def add_one(req: AddRequest, request: Request):
|
|
rows = watchlist.add(req.symbol, req.note)
|
|
return {"symbols": _with_names(rows, request)}
|
|
|
|
|
|
@router.post("/batch")
|
|
def add_batch(req: BatchAddRequest, request: Request):
|
|
for sym in req.symbols:
|
|
watchlist.add(sym, req.note)
|
|
return {"symbols": _with_names(watchlist.list_symbols(), request), "added": len(req.symbols)}
|
|
|
|
|
|
@router.post("/{symbol}/top")
|
|
def move_one_to_top(symbol: str, request: Request):
|
|
rows = watchlist.move_to_top(symbol)
|
|
return {"symbols": _with_names(rows, request)}
|
|
|
|
|
|
@router.delete("/{symbol}")
|
|
def remove_one(symbol: str, request: Request):
|
|
rows = watchlist.remove(symbol)
|
|
return {"symbols": _with_names(rows, request)}
|
|
|
|
|
|
@router.delete("")
|
|
def clear_all():
|
|
"""清空自选列表。"""
|
|
count = watchlist.clear()
|
|
return {"removed": count}
|
|
|
|
|
|
# 自选页需要的列
|
|
_WATCHLIST_COLS = [
|
|
"symbol", "close", "change_pct", "change_amount", "amount",
|
|
"turnover_rate",
|
|
"amplitude", "annual_vol_20d",
|
|
"vol_ratio_5d",
|
|
"ma5", "ma10", "ma20", "ma60",
|
|
"vol_ma5", "vol_ma10",
|
|
"high_60d", "low_60d",
|
|
"rsi_6", "rsi_14", "rsi_24",
|
|
"macd_dif", "macd_dea", "macd_hist",
|
|
"kdj_k", "kdj_d", "kdj_j",
|
|
"boll_upper", "boll_lower",
|
|
"atr_14",
|
|
"momentum_5d", "momentum_10d", "momentum_20d", "momentum_30d", "momentum_60d",
|
|
"consecutive_limit_ups", "consecutive_limit_downs",
|
|
"signal_limit_up", "signal_limit_down", "signal_volume_surge",
|
|
"signal_ma_golden_5_20", "signal_macd_golden", "signal_n_day_high",
|
|
"signal_boll_breakout_upper", "signal_ma20_breakout",
|
|
"signal_ma_dead_5_20", "signal_macd_dead", "signal_n_day_low",
|
|
"signal_boll_breakdown_lower", "signal_ma20_breakdown",
|
|
]
|
|
|
|
|
|
@router.get("/enriched")
|
|
def watchlist_enriched(
|
|
request: Request,
|
|
ext_columns: str | None = Query(None, description="逗号分隔的 ext 列: config_id.field_name"),
|
|
):
|
|
"""自选股 enriched 数据 — 直接从 enriched 最新日读取, 无即时计算。
|
|
|
|
ext_columns 参数示例: "industry_rating.score,fund_flow.net_inflow"
|
|
会动态 LEFT JOIN 对应的 ext_{config_id} DuckDB view。
|
|
"""
|
|
t0 = time.perf_counter()
|
|
|
|
repo = request.app.state.repo
|
|
symbols = [r["symbol"] for r in watchlist.list_symbols()]
|
|
if not symbols:
|
|
return {"rows": [], "as_of": None, "elapsed_ms": 0}
|
|
|
|
df_e, cache_date = repo.get_enriched_latest()
|
|
if df_e.is_empty():
|
|
return {"rows": [], "as_of": None, "elapsed_ms": 0}
|
|
|
|
# 按 symbol 过滤
|
|
df = df_e.filter(pl.col("symbol").is_in(symbols))
|
|
if df.is_empty():
|
|
return {"rows": [], "as_of": str(cache_date) if cache_date else None, "elapsed_ms": 0}
|
|
|
|
# JOIN instruments 取 name + float_shares
|
|
df_i = repo.get_instruments()
|
|
if not df_i.is_empty() and "name" in df_i.columns:
|
|
inst_cols = [c for c in ["symbol", "name", "float_shares"] if c in df_i.columns]
|
|
df = df.join(df_i.select(inst_cols), on="symbol", how="left")
|
|
|
|
# 选择内置需要的列
|
|
keep = [c for c in _WATCHLIST_COLS + ["name", "float_shares"] if c in df.columns]
|
|
df = df.select(keep)
|
|
|
|
# 动态 JOIN 扩展数据表
|
|
ext_specs = _parse_ext_columns(ext_columns) if ext_columns else []
|
|
if ext_specs:
|
|
db = repo.store.db
|
|
data_dir = repo.store.data_dir
|
|
from app.services.ext_data import ExtConfigStore
|
|
from app.api.ext_data import _read_ext_dataframe
|
|
|
|
ext_store = ExtConfigStore(data_dir)
|
|
configs = {c.id: c for c in ext_store.load_all()}
|
|
|
|
for config_id, field_name in ext_specs:
|
|
view_name = f"ext_{config_id}"
|
|
ext_col_name = f"{config_id}__{field_name}"
|
|
try:
|
|
# 扩展时序数据必须只取最新分区；否则一个 symbol 会按历史分区数被 JOIN 放大。
|
|
cfg = configs.get(config_id)
|
|
if cfg:
|
|
ext_df, _ = _read_ext_dataframe(cfg, data_dir)
|
|
else:
|
|
ext_df = pl.from_arrow(db.query(
|
|
f"SELECT symbol, \"{field_name}\" FROM {view_name}"
|
|
).arrow())
|
|
if not ext_df.is_empty() and "symbol" in ext_df.columns:
|
|
ext_df = (
|
|
ext_df
|
|
.select(["symbol", field_name])
|
|
.unique(subset=["symbol"], keep="last")
|
|
.rename({field_name: ext_col_name})
|
|
)
|
|
df = df.join(ext_df.select(["symbol", ext_col_name]), on="symbol", how="left")
|
|
except Exception:
|
|
# view 不存在或字段不存在，尝试直接读 parquet
|
|
cfg = configs.get(config_id)
|
|
if cfg:
|
|
try:
|
|
ext_df, _ = _read_ext_dataframe(cfg, data_dir)
|
|
if not ext_df.is_empty() and "symbol" in ext_df.columns and field_name in ext_df.columns:
|
|
ext_df = (
|
|
ext_df
|
|
.select(["symbol", field_name])
|
|
.unique(subset=["symbol"], keep="last")
|
|
.rename({field_name: ext_col_name})
|
|
)
|
|
df = df.join(ext_df, on="symbol", how="left")
|
|
except Exception as e2:
|
|
logger.debug("ext join fallback failed for %s.%s: %s", config_id, field_name, e2)
|
|
|
|
# sanitize NaN / Inf
|
|
float_cols = [c for c in df.columns if df[c].dtype.is_float()]
|
|
if float_cols:
|
|
df = df.with_columns([
|
|
pl.when(pl.col(c).is_nan() | pl.col(c).is_infinite())
|
|
.then(None)
|
|
.otherwise(pl.col(c))
|
|
.alias(c)
|
|
for c in float_cols
|
|
])
|
|
|
|
# 按自选添加顺序（新加的在前）重排行
|
|
order_map = {s: i for i, s in enumerate(symbols)}
|
|
df = df.with_columns(pl.col("symbol").map_elements(lambda s: order_map.get(s, len(symbols)), return_dtype=pl.Int32).alias("_sort_order"))
|
|
df = df.sort("_sort_order").drop("_sort_order")
|
|
|
|
rows = df.to_dicts()
|
|
elapsed = (time.perf_counter() - t0) * 1000
|
|
return {"rows": rows, "as_of": str(cache_date) if cache_date else None, "elapsed_ms": elapsed}
|
|
|
|
|
|
def _parse_ext_columns(ext_columns: str) -> list[tuple[str, str]]:
|
|
"""解析 'config_id1.field1,config_id2.field2' 为 [(config_id, field_name), ...]"""
|
|
result = []
|
|
for part in ext_columns.split(","):
|
|
part = part.strip()
|
|
if "." not in part:
|
|
continue
|
|
config_id, field_name = part.split(".", 1)
|
|
config_id = config_id.strip()
|
|
field_name = field_name.strip()
|
|
if config_id and field_name:
|
|
result.append((config_id, field_name))
|
|
return result
|