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]: hub = getattr(self, "try_index_quotes", None) if callable(hub): rows = hub() if rows: try: return self._hub_realtime_market_indices(requested_date, rows) except TushareError: pass try: payload = self._tushare_realtime_market_indices(requested_date) marker = getattr(self, "record_datahub_legacy", None) if callable(marker): marker("index_quotes", "tushare_rt_idx_k") return payload except TushareError: payload = self._free_realtime_market_indices(requested_date) marker = getattr(self, "record_datahub_legacy", None) if callable(marker): marker("index_quotes", str(payload.get("source") or "eastmoney_push2")) return payload def _hub_realtime_market_indices( self, requested_date: str, rows: list[dict[str, Any]], ) -> dict[str, Any]: trade_date, _ = self.resolve_trade_context(requested_date) index_names = { "000001.SH": "上证指数", "399001.SZ": "深证成指", "399006.SZ": "创业板指", } by_code = {str(row.get("ts_code") or ""): row for row in rows} by_symbol = {str(row.get("code") or ""): row for row in rows} indices = [] for ts_code, name in index_names.items(): row = by_code.get(ts_code) or by_symbol.get(ts_code.split(".")[0]) if not row: continue close = _number(row.get("price") if row.get("price") not in (None, "") else row.get("close")) previous_close = _number( row.get("previous_close") if row.get("previous_close") not in (None, "") else row.get("pre_close") ) if close <= 0 or previous_close <= 0: continue amount = _number(row.get("amount")) amount_billion = _number(row.get("amount_billion")) if not amount_billion and amount: amount_billion = round(amount / 100_000_000, 2) indices.append( { "ts_code": ts_code, "name": str(row.get("name") or name).strip(), "trade_date": trade_date, "close": close, "pct_chg": round( _number(row.get("pct_chg")) or (close / previous_close - 1) * 100, 3, ), "return_5d": 0, "amount_billion": amount_billion, "quote_time": str(row.get("quote_time") or ""), "source": "datahub", } ) if len(indices) != 3: raise TushareError("Realtime index quotes are incomplete") return { "trade_date": trade_date, "source": "datahub", "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": 0, "average_return_20d": 0, }, } def _tushare_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, }, } def _free_realtime_market_indices(self, requested_date: str) -> dict[str, Any]: trade_date, _ = self.resolve_trade_context(requested_date) aggregator = getattr(self, "realtime_aggregator", None) if aggregator is None: raise TushareError("免费实时源未配置") quotes = aggregator.eastmoney_indices() index_names = { "000001": ("000001.SH", "上证指数"), "399001": ("399001.SZ", "深证成指"), "399006": ("399006.SZ", "创业板指"), } indices = [] for quote in quotes: mapped = index_names.get(str(quote.get("code") or "")) if not mapped: continue ts_code, name = mapped close = _number(quote.get("price")) previous_close = _number(quote.get("previous_close")) if close <= 0 or previous_close <= 0: continue indices.append( { "ts_code": ts_code, "name": str(quote.get("name") or name).strip(), "trade_date": trade_date, "close": close, "pct_chg": round(_number(quote.get("change")) or (close / previous_close - 1) * 100, 3), "return_5d": 0, "amount_billion": round(_number(quote.get("amount_billion")), 2), "quote_time": quote.get("quote_time") or "", "source": quote.get("source") or "eastmoney_push2", } ) if len(indices) != 3: raise TushareError("Realtime index quotes are incomplete") return { "trade_date": trade_date, "source": "eastmoney_push2", "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": 0, "average_return_20d": 0, }, }