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_dragon_tiger,
    normalize_index_daily,
    normalize_limit_event,
    normalize_moneyflow,
    normalize_popularity,
    normalize_sector_daily,
    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",
    "limit_list_d": (
        "trade_date,ts_code,industry,name,close,pct_chg,amount,limit_amount,"
        "float_mv,total_mv,turnover_ratio,fd_amount,first_time,last_time,"
        "open_times,up_stat,limit_times,limit_type"
    ),
    "ths_hot": "ts_code,ts_name,hot,rank,pct_change,current_price,concept,data_type,trade_date",
    "dc_hot": "ts_code,ts_name,rank,pct_change,current_price,hot,concept,data_type,trade_date",
    "hm_detail": "trade_date,ts_code,ts_name,buy_amount,sell_amount,net_amount,hm_name,hm_orgs,tag",
    "hm_list": "name,desc,orgs",
    "top_list": "trade_date,ts_code,name,pct_change,reason",
    "top_inst": "trade_date,ts_code,exalter,buy,buy_rate,sell,sell_rate,net_buy,side,reason",
    "ths_index": "ts_code,name,count,exchange,list_date,type",
    "ths_daily": "ts_code,trade_date,open,high,low,close,pre_close,pct_change,vol,turnover_rate",
    "dc_index": "ts_code,trade_date,name,open,high,low,close,pre_close,pct_change,vol,amount,turnover_rate",
    "sw_daily": "ts_code,trade_date,name,open,high,low,close,pct_change,vol,amount",
}

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",
    "limit_events": "limit_list_d",
    "popularity": "ths_hot",
    "dragon_tiger": "hm_detail",
    "sector_daily": "ths_daily",
}

WEBSITE_INDEX_CODES = ("000001.SH", "399001.SZ", "399006.SZ", "000300.SH")
DEFAULT_INDEX_CODES = WEBSITE_INDEX_CODES
LIMIT_TYPES = ("U", "D", "Z")


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]]:
        if dataset == "limit_events":
            return self.fetch_limit_events(str(params.get("trade_date") or ""))
        if dataset == "popularity":
            return self.fetch_popularity(str(params.get("trade_date") or ""))
        if dataset == "dragon_tiger":
            return self.fetch_dragon_tiger(str(params.get("trade_date") or ""))
        if dataset == "sector_daily":
            return self.fetch_sector_daily(str(params.get("trade_date") or ""))
        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:
            query_params.setdefault("ts_code", DEFAULT_INDEX_CODES[0])
        return self._query(api_name, query_params, fields)

    def fetch_limit_events(self, trade_date: str) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for limit_type in LIMIT_TYPES:
            part = self._query(
                "limit_list_d",
                {"trade_date": trade_date, "limit_type": limit_type},
                TUSHARE_FIELDS["limit_list_d"],
            )
            for row in part:
                row = dict(row)
                row.setdefault("limit_type", limit_type)
                rows.append(row)
        return rows

    def fetch_popularity(self, trade_date: str) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for api_name, source in (("ths_hot", "ths"), ("dc_hot", "dc")):
            for row in self._query(api_name, {"trade_date": trade_date}, TUSHARE_FIELDS[api_name]):
                item = dict(row)
                item["source"] = source
                item.setdefault("trade_date", trade_date)
                rows.append(item)
        return rows

    def fetch_dragon_tiger(self, trade_date: str) -> list[dict[str, Any]]:
        details = self._query("hm_detail", {"trade_date": trade_date}, TUSHARE_FIELDS["hm_detail"])
        top_rows = self._query("top_list", {"trade_date": trade_date}, TUSHARE_FIELDS["top_list"])
        context = {
            str(row.get("ts_code") or ""): row
            for row in top_rows
            if str(row.get("ts_code") or "")
        }
        rows: list[dict[str, Any]] = []
        for row in details:
            item = dict(row)
            stock = context.get(str(item.get("ts_code") or ""), {})
            if item.get("pct_change") is None and stock.get("pct_change") is not None:
                item["pct_change"] = stock.get("pct_change")
            if not item.get("reason") and stock.get("reason"):
                item["reason"] = stock.get("reason")
            if not item.get("ts_name") and stock.get("name"):
                item["ts_name"] = stock.get("name")
            rows.append(item)
        return rows

    def fetch_sector_daily(self, trade_date: str) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for api_name, family in (("ths_daily", "ths"), ("dc_index", "dc"), ("sw_daily", "sw")):
            try:
                part = self._query(api_name, {"trade_date": trade_date}, TUSHARE_FIELDS[api_name])
            except AdapterError:
                part = []
            for row in part:
                item = dict(row)
                item["family"] = family
                rows.append(item)
        return rows

    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]]:
        if dataset in {"limit_events", "limit_list_d"}:
            return [normalize_limit_event(row) for row in rows]
        if dataset == "popularity":
            return [normalize_popularity(row, source=str(row.get("source") or "")) for row in rows]
        if dataset == "dragon_tiger":
            return [normalize_dragon_tiger(row) for row in rows]
        if dataset == "sector_daily":
            return [
                normalize_sector_daily(row, family=str(row.get("family") or "ths"))
                for row in rows
            ]
        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 (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
            raise AdapterError(f"Tushare 请求失败: {exc}") from exc
        if result.get("code") not in (0, "0", None):
            raise AdapterError(str(result.get("msg") or f"Tushare error {result.get('code')}"))
        data = result.get("data") or {}
        items = data.get("items") or []
        fields_list = data.get("fields") or (fields.split(",") if fields else [])
        return [dict(zip(fields_list, item)) for item in items]