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]