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

from __future__ import annotations

import json
import time
import urllib.error
import urllib.request
from typing import Any, Callable

from datahub.adapters.base import AdapterError, MarketAdapter
from datahub.normalize import (
    normalize_auction,
    normalize_calendar,
    normalize_daily,
    normalize_index_daily,
    normalize_moneyflow,
    normalize_stock,
    normalize_valuation,
)

TUSHARE_URL = "http://api.tushare.pro"

TUSHARE_FIELDS = {
    "trade_cal": "exchange,cal_date,is_open,pretrade_date",
    "stock_basic": "ts_code,symbol,name,area,industry,market,list_status,list_date",
    "daily": "ts_code,trade_date,open,high,low,close,pct_chg,vol,amount",
    "daily_basic": "ts_code,trade_date,turnover_rate,volume_ratio,total_mv,circ_mv,pe_ttm,pb,ps_ttm,dv_ttm",
    "adj_factor": "ts_code,trade_date,adj_factor",
    "index_daily": "ts_code,trade_date,open,high,low,close,pct_chg,vol,amount",
    "moneyflow": (
        "ts_code,trade_date,buy_sm_amount,sell_sm_amount,buy_md_amount,sell_md_amount,"
        "buy_lg_amount,sell_lg_amount,buy_elg_amount,sell_elg_amount,net_mf_amount"
    ),
    "stk_auction": "ts_code,trade_date,vol,price,amount,pre_close,turnover_rate,volume_ratio,float_share",
}

DATASET_API = {
    "calendar": "trade_cal",
    "stocks": "stock_basic",
    "daily": "daily",
    "valuation": "daily_basic",
    "adj_factor": "adj_factor",
    "index_daily": "index_daily",
    "moneyflow": "moneyflow",
    "auction": "stk_auction",
}

DEFAULT_INDEX_CODES = ("000001.SH", "399001.SZ", "399006.SZ", "000300.SH")


class TushareAdapter(MarketAdapter):
    name = "tushare"

    def __init__(
        self,
        token: str,
        timeout: int = 30,
        transport: Callable[[str, dict[str, Any], str], list[dict[str, Any]]] | None = None,
    ) -> None:
        self.token = token
        self.timeout = timeout
        self._transport = transport

    def probe(self) -> dict[str, Any]:
        if not self.token:
            return {"provider": self.name, "configured": False, "state": "unconfigured"}
        started = time.perf_counter()
        try:
            rows = self.fetch("calendar", {"exchange": "SSE", "start_date": "20200102", "end_date": "20200102"})
        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": "ok" if rows else "empty",
            "latency_ms": round((time.perf_counter() - started) * 1000),
        }

    def fetch(self, dataset: str, params: dict[str, Any]) -> list[dict[str, Any]]:
        api_name = DATASET_API.get(dataset, dataset)
        fields = TUSHARE_FIELDS.get(api_name, "")
        query_params = dict(params)
        if api_name == "stock_basic" and "list_status" not in query_params:
            query_params["list_status"] = "L"
        if api_name == "trade_cal" and "exchange" not in query_params:
            query_params["exchange"] = "SSE"
        if api_name == "index_daily" and "ts_code" not in query_params:
            # Caller typically loops codes; a missing code would pull nothing useful.
            query_params.setdefault("ts_code", DEFAULT_INDEX_CODES[0])
        return self._query(api_name, query_params, fields)

    def fetch_index_daily(self, trade_date: str, codes: tuple[str, ...] = DEFAULT_INDEX_CODES) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for ts_code in codes:
            rows.extend(self.fetch("index_daily", {"ts_code": ts_code, "trade_date": trade_date}))
        return rows

    def normalize(self, dataset: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        mapping = {
            "calendar": normalize_calendar,
            "trade_cal": normalize_calendar,
            "stocks": normalize_stock,
            "stock_basic": normalize_stock,
            "daily": normalize_daily,
            "valuation": normalize_valuation,
            "daily_basic": normalize_valuation,
            "moneyflow": normalize_moneyflow,
            "auction": normalize_auction,
            "stk_auction": normalize_auction,
            "index_daily": normalize_index_daily,
        }
        fn = mapping.get(dataset)
        if fn is None:
            if dataset == "adj_factor":
                return [
                    {
                        "ts_code": str(row.get("ts_code") or "").upper(),
                        "trade_date": str(row.get("trade_date") or ""),
                        "adj_factor": row.get("adj_factor"),
                    }
                    for row in rows
                ]
            raise AdapterError(f"unsupported dataset: {dataset}")
        return [fn(row) for row in rows]

    def _query(self, api_name: str, params: dict[str, Any], fields: str) -> list[dict[str, Any]]:
        if self._transport is not None:
            return self._transport(api_name, params, fields)
        if not self.token:
            raise AdapterError("Tushare token 未配置")
        payload = json.dumps(
            {"api_name": api_name, "token": self.token, "params": params, "fields": fields}
        ).encode("utf-8")
        request = urllib.request.Request(
            TUSHARE_URL,
            data=payload,
            headers={"Content-Type": "application/json", "User-Agent": "XiaobaiDatahub/0.1"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                result = json.loads(response.read().decode("utf-8"))
        except json.JSONDecodeError:
            raise AdapterError("Tushare returned invalid json") from None
        except (urllib.error.URLError, TimeoutError) as exc:
            raise AdapterError(f"Tushare request failed: {exc}") from exc
        if result.get("code") != 0:
            raise AdapterError(result.get("msg") or "Tushare returned an unknown error")
        data = result.get("data") or {}
        columns = data.get("fields") or []
        return [dict(zip(columns, item)) for item in data.get("items") or []]