xiaobai-review/backend/data/providers/tushare_indices.py

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,
            },
        }