xiaobai-review/xiaobai-datahub/datahub/adapters/eastmoney.py

from __future__ import annotations

import json
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime
from typing import Any

from datahub.adapters.base import AdapterError, MarketAdapter
from datahub.numbers import finite_number, round4

EASTMONEY_INDEX_URL = "https://push2.eastmoney.com/api/qt/ulist.np/get"
EASTMONEY_CLIST_URL = "https://push2.eastmoney.com/api/qt/clist/get"
TRENDS_URL = "https://push2delay.eastmoney.com/api/qt/stock/trends2/get"
BROWSER_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"
)
INDEX_SECIDS = {
    "000001.SH": "1.000001",
    "399001.SZ": "0.399001",
    "399006.SZ": "0.399006",
}


class EastmoneyAdapter(MarketAdapter):
    name = "eastmoney"

    def __init__(self, timeout: int = 8) -> None:
        self.timeout = timeout

    def probe(self) -> dict[str, Any]:
        started = time.perf_counter()
        try:
            rows = self.fetch_indices()
            state = "ok" if len(rows) == 3 else "empty"
        except AdapterError as exc:
            return {
                "provider": self.name,
                "configured": True,
                "state": "error",
                "message": str(exc),
                "latency_ms": round((time.perf_counter() - started) * 1000),
            }
        return {
            "provider": self.name,
            "configured": True,
            "state": state,
            "latency_ms": round((time.perf_counter() - started) * 1000),
        }

    def fetch(self, dataset: str, params: dict[str, Any]) -> list[dict[str, Any]]:
        if dataset in {"indexes_quotes", "index_quotes"}:
            return self.fetch_indices()
        if dataset in {"quotes", "quotes_latest"}:
            codes = params.get("codes") or []
            if isinstance(codes, str):
                codes = [item.strip() for item in codes.split(",") if item.strip()]
            return self.fetch_quotes(list(codes))
        raise AdapterError(f"{self.name} unsupported dataset: {dataset}")

    def normalize(self, dataset: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        return list(rows)

    def fetch_indices(self) -> list[dict[str, Any]]:
        payload = self._get_json(
            EASTMONEY_INDEX_URL,
            {
                "secids": "1.000001,0.399001,0.399006",
                "fltt": "2",
                "invt": "2",
                "fields": "f12,f14,f2,f3,f4,f15,f16,f17,f18,f6,f124",
            },
            referer="https://quote.eastmoney.com/",
        )
        rows = list((payload.get("data") or {}).get("diff") or [])
        result = []
        for row in rows:
            code = str(row.get("f12") or "")
            if code not in {"000001", "399001", "399006"}:
                continue
            epoch = int(finite_number(row.get("f124")) or 0)
            ts_code = f"{code}.SH" if code.startswith("0") and code == "000001" else f"{code}.SZ"
            if code == "000001":
                ts_code = "000001.SH"
            result.append(
                {
                    "ts_code": ts_code,
                    "code": code,
                    "name": row.get("f14") or code,
                    "price": round4(finite_number(row.get("f2"))),
                    "pct_chg": round4(finite_number(row.get("f3"))),
                    "change_amount": round4(finite_number(row.get("f4"))),
                    "open": round4(finite_number(row.get("f17"))),
                    "high": round4(finite_number(row.get("f15"))),
                    "low": round4(finite_number(row.get("f16"))),
                    "previous_close": round4(finite_number(row.get("f18"))),
                    "amount": round4(finite_number(row.get("f6"))),
                    "quote_time_epoch": epoch,
                    "quote_time": (
                        datetime.fromtimestamp(epoch).astimezone().isoformat(timespec="seconds")
                        if epoch
                        else ""
                    ),
                    "source": "eastmoney_push2",
                }
            )
        if len(result) != 3:
            raise AdapterError(f"Eastmoney returned {len(result)}/3 indices")
        return result

    def fetch_quotes(self, codes: list[str]) -> list[dict[str, Any]]:
        # Eastmoney clist does not accept arbitrary code lists well; use ulist.np for batches.
        secids = []
        for code in codes:
            ts = str(code or "").upper()
            symbol = ts.split(".")[0]
            if ts.endswith(".SH") or symbol.startswith(("5", "6", "9")):
                secids.append(f"1.{symbol}")
            else:
                secids.append(f"0.{symbol}")
        if not secids:
            return []
        payload = self._get_json(
            EASTMONEY_INDEX_URL,
            {
                "secids": ",".join(secids[:60]),
                "fltt": "2",
                "invt": "2",
                "fields": "f12,f14,f2,f3,f4,f15,f16,f17,f18,f5,f6,f8,f124",
            },
            referer="https://quote.eastmoney.com/",
        )
        rows = list((payload.get("data") or {}).get("diff") or [])
        result = []
        for row in rows:
            symbol = str(row.get("f12") or "")
            if not symbol:
                continue
            ts_code = f"{symbol}.SH" if symbol.startswith(("5", "6", "9")) else f"{symbol}.SZ"
            epoch = int(finite_number(row.get("f124")) or 0)
            result.append(
                {
                    "ts_code": ts_code,
                    "name": row.get("f14") or symbol,
                    "price": round4(finite_number(row.get("f2"))),
                    "pct_chg": round4(finite_number(row.get("f3"))),
                    "change_amount": round4(finite_number(row.get("f4"))),
                    "open": round4(finite_number(row.get("f17"))),
                    "high": round4(finite_number(row.get("f15"))),
                    "low": round4(finite_number(row.get("f16"))),
                    "previous_close": round4(finite_number(row.get("f18"))),
                    "volume": round4(finite_number(row.get("f5"))),
                    "amount": round4(finite_number(row.get("f6"))),
                    "turnover_rate": round4(finite_number(row.get("f8"))),
                    "quote_time_epoch": epoch,
                    "quote_time": (
                        datetime.fromtimestamp(epoch).astimezone().isoformat(timespec="seconds")
                        if epoch
                        else ""
                    ),
                    "source": "eastmoney_push2",
                }
            )
        return result

    def fetch_intraday(self, ts_code: str) -> dict[str, Any]:
        code = str(ts_code or "").upper()
        if code in INDEX_SECIDS:
            secid = INDEX_SECIDS[code]
            entity = "index"
            identifier = code
        else:
            symbol = code.split(".")[0]
            market = "1" if symbol.startswith(("5", "6", "9")) else "0"
            secid = f"{market}.{symbol}"
            entity = "stock"
            identifier = symbol
        payload = self._get_json(
            TRENDS_URL,
            {
                "secid": secid,
                "fields1": "f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13",
                "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
                "iscr": "0",
                "ndays": "1",
            },
            referer="https://quote.eastmoney.com/",
        )
        data = payload.get("data") or {}
        points = []
        for raw in data.get("trends") or []:
            point = _parse_trend(raw)
            if point:
                points.append(point)
        if not points:
            raise AdapterError("No intraday chart data returned")
        return {
            "entity_type": entity,
            "identifier": identifier,
            "ts_code": code if "." in code else f"{identifier}.{'SH' if identifier.startswith(('5','6','9')) else 'SZ'}",
            "name": str(data.get("name") or ""),
            "code": str(data.get("code") or identifier),
            "trade_date": points[-1]["date"],
            "previous_close": round4(finite_number(data.get("preClose"))),
            "points": points,
            "source": "eastmoney_trends2",
        }

    def _get_json(self, url: str, params: dict[str, str], referer: str) -> dict[str, Any]:
        request_url = f"{url}?{urllib.parse.urlencode(params)}"
        request = urllib.request.Request(
            request_url,
            headers={
                "Accept": "application/json,text/plain,*/*",
                "User-Agent": BROWSER_UA,
                "Referer": referer,
            },
            method="GET",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                return json.loads(response.read().decode("utf-8"))
        except Exception as exc:
            raise AdapterError(f"eastmoney request failed: {exc}") from exc


def _parse_trend(raw: Any) -> dict[str, Any] | None:
    text = str(raw or "")
    parts = text.split(",")
    if len(parts) < 8:
        return None
    stamp = parts[0]
    try:
        when = datetime.strptime(stamp, "%Y-%m-%d %H:%M")
    except ValueError:
        return None
    return {
        "time": when.strftime("%H:%M"),
        "date": when.strftime("%Y-%m-%d"),
        "open": round4(finite_number(parts[1])),
        "close": round4(finite_number(parts[2])),
        "high": round4(finite_number(parts[3])),
        "low": round4(finite_number(parts[4])),
        "avg_price": round4(finite_number(parts[7] if len(parts) > 7 else parts[2])),
        "volume": round4(finite_number(parts[5])),
        "amount": round4(finite_number(parts[6])),
    }