119 lines
5.3 KiB
Python
119 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta
|
|
from typing import Any
|
|
|
|
from backend.data.numbers import finite_number as _number
|
|
from backend.data.providers.tushare_transport import TushareError
|
|
|
|
|
|
class IndexMixin:
|
|
def market_indices(self, requested_date: str, lookback_days: int = 45) -> dict[str, Any]:
|
|
trade_date, _ = self.resolve_trade_context(requested_date)
|
|
end = datetime.strptime(trade_date, "%Y%m%d")
|
|
start_date = (end - timedelta(days=max(30, lookback_days * 2))).strftime("%Y%m%d")
|
|
index_names = {
|
|
"000001.SH": "上证指数",
|
|
"399001.SZ": "深证成指",
|
|
"399006.SZ": "创业板指",
|
|
}
|
|
indices = []
|
|
for ts_code, name in index_names.items():
|
|
rows = self.query(
|
|
"index_daily",
|
|
{"ts_code": ts_code, "start_date": start_date, "end_date": trade_date},
|
|
"ts_code,trade_date,close,pct_chg,vol,amount",
|
|
)
|
|
rows.sort(key=lambda item: str(item.get("trade_date") or ""))
|
|
if not rows:
|
|
continue
|
|
latest = rows[-1]
|
|
close = _number(latest.get("close"))
|
|
close_5d = _number(rows[-6].get("close")) if len(rows) >= 6 else _number(rows[0].get("close"))
|
|
close_20d = _number(rows[-21].get("close")) if len(rows) >= 21 else _number(rows[0].get("close"))
|
|
indices.append(
|
|
{
|
|
"ts_code": ts_code,
|
|
"name": name,
|
|
"trade_date": str(latest.get("trade_date") or trade_date),
|
|
"close": close,
|
|
"pct_chg": round(_number(latest.get("pct_chg")), 3),
|
|
"return_5d": round((close / close_5d - 1) * 100, 3) if close_5d else 0,
|
|
"return_20d": round((close / close_20d - 1) * 100, 3) if close_20d else 0,
|
|
"amount_billion": round(_number(latest.get("amount")) / 100000, 2),
|
|
}
|
|
)
|
|
if not indices:
|
|
raise TushareError(f"No index data returned for {trade_date}")
|
|
return {
|
|
"trade_date": trade_date,
|
|
"source": "tushare",
|
|
"realtime": False,
|
|
"precise": all(item["trade_date"] == trade_date for item in indices),
|
|
"indices": indices,
|
|
"aggregate": {
|
|
"average_pct_chg": round(sum(item["pct_chg"] for item in indices) / len(indices), 3),
|
|
"average_return_5d": round(sum(item["return_5d"] for item in indices) / len(indices), 3),
|
|
"average_return_20d": round(sum(item["return_20d"] for item in indices) / len(indices), 3),
|
|
},
|
|
}
|
|
|
|
def realtime_market_indices(self, requested_date: str) -> dict[str, Any]:
|
|
trade_date, _ = self.resolve_trade_context(requested_date)
|
|
index_names = {
|
|
"000001.SH": "上证指数",
|
|
"399001.SZ": "深证成指",
|
|
"399006.SZ": "创业板指",
|
|
}
|
|
rows = self.query("rt_idx_k", {"ts_code": ",".join(index_names)}, "")
|
|
row_map = {str(row.get("ts_code") or ""): row for row in rows}
|
|
indices = []
|
|
for ts_code, name in index_names.items():
|
|
row = row_map.get(ts_code)
|
|
if not row:
|
|
continue
|
|
close = _number(row.get("close"))
|
|
previous_close = _number(row.get("pre_close"))
|
|
if close <= 0 or previous_close <= 0:
|
|
continue
|
|
history = self.query(
|
|
"index_daily",
|
|
{
|
|
"ts_code": ts_code,
|
|
"start_date": (datetime.strptime(trade_date, "%Y%m%d") - timedelta(days=20)).strftime("%Y%m%d"),
|
|
"end_date": trade_date,
|
|
},
|
|
"ts_code,trade_date,close,pct_chg",
|
|
)
|
|
history.sort(key=lambda item: str(item.get("trade_date") or ""))
|
|
previous_closes = [
|
|
_number(item.get("close")) for item in history
|
|
if str(item.get("trade_date") or "") < trade_date and _number(item.get("close")) > 0
|
|
]
|
|
close_5d = previous_closes[-5] if len(previous_closes) >= 5 else previous_closes[0] if previous_closes else previous_close
|
|
indices.append(
|
|
{
|
|
"ts_code": ts_code,
|
|
"name": str(row.get("name") or name).strip(),
|
|
"trade_date": trade_date,
|
|
"close": close,
|
|
"pct_chg": round((close / previous_close - 1) * 100, 3),
|
|
"return_5d": round((close / close_5d - 1) * 100, 3) if close_5d else 0,
|
|
"amount_billion": round(_number(row.get("amount")) / 100000000, 2),
|
|
}
|
|
)
|
|
if len(indices) != len(index_names):
|
|
raise TushareError("Realtime index quotes are incomplete")
|
|
return {
|
|
"trade_date": trade_date,
|
|
"source": "tushare_rt_idx_k",
|
|
"realtime": True,
|
|
"precise": True,
|
|
"indices": indices,
|
|
"aggregate": {
|
|
"average_pct_chg": round(sum(item["pct_chg"] for item in indices) / len(indices), 3),
|
|
"average_return_5d": round(sum(item["return_5d"] for item in indices) / len(indices), 3),
|
|
"average_return_20d": 0,
|
|
},
|
|
}
|