290 lines
10 KiB
Python
290 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import urllib.error
|
|
import urllib.request
|
|
from collections.abc import Callable
|
|
from dataclasses import replace
|
|
from datetime import datetime, timedelta
|
|
from typing import Any
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from backend.data.contracts import (
|
|
DataSource,
|
|
DataUsage,
|
|
ObservationMetadata,
|
|
ProviderResult,
|
|
SnapshotState,
|
|
)
|
|
from backend.data.providers.base import ProviderError
|
|
|
|
SHANGHAI = ZoneInfo("Asia/Shanghai")
|
|
|
|
|
|
class TushareProvider:
|
|
source = DataSource.TUSHARE
|
|
url = "http://api.tushare.pro"
|
|
|
|
def __init__(self, token: str | None | Callable[[], str | None], timeout: int = 20) -> None:
|
|
self._token_provider = token if callable(token) else lambda: token
|
|
self._timeout = timeout
|
|
|
|
@property
|
|
def configured(self) -> bool:
|
|
return bool(self._token())
|
|
|
|
def calendar(self, start_date: str, end_date: str) -> ProviderResult:
|
|
result = self._query(
|
|
"trade_cal",
|
|
{"exchange": "SSE", "start_date": _compact(start_date), "end_date": _compact(end_date)},
|
|
"cal_date,is_open,pretrade_date",
|
|
unit="calendar_day",
|
|
)
|
|
days = (
|
|
datetime.fromisoformat(end_date).date() - datetime.fromisoformat(start_date).date()
|
|
).days + 1
|
|
return ProviderResult(
|
|
result.rows,
|
|
replace(result.metadata, coverage=min(len(result.rows) / max(days, 1), 1)),
|
|
)
|
|
|
|
def entities(self) -> ProviderResult:
|
|
rows: list[dict[str, Any]] = []
|
|
for status in ("L", "P", "D"):
|
|
result = self._query(
|
|
"stock_basic",
|
|
{"exchange": "", "list_status": status},
|
|
"ts_code,symbol,name,industry,list_status,list_date,delist_date",
|
|
unit="entity",
|
|
)
|
|
rows.extend(result.rows)
|
|
return ProviderResult(
|
|
tuple(rows), _metadata(self.source, "entity", min(len(rows) / 5300, 1))
|
|
)
|
|
|
|
def daily(self, entity_type: str, identifier: str, end_date: str) -> ProviderResult:
|
|
api_name = "index_daily" if entity_type == "index" else "daily"
|
|
if entity_type in {"sector", "theme"}:
|
|
api_name = "ths_daily"
|
|
end = datetime.strptime(_compact(end_date), "%Y%m%d")
|
|
start = (end - timedelta(days=380)).strftime("%Y%m%d")
|
|
return self._query(
|
|
api_name,
|
|
{"ts_code": identifier, "start_date": start, "end_date": end.strftime("%Y%m%d")},
|
|
"ts_code,trade_date,open,high,low,close,vol,amount,pct_chg",
|
|
unit="yuan/share",
|
|
adjustment="unadjusted",
|
|
)
|
|
|
|
def minute(self, entity_type: str, identifier: str, trade_date: str) -> ProviderResult:
|
|
if entity_type != "stock":
|
|
raise ProviderError("Tushare minute charts only support stocks")
|
|
date = _display(trade_date)
|
|
return self._query(
|
|
"stk_mins",
|
|
{
|
|
"ts_code": identifier,
|
|
"freq": "1min",
|
|
"start_date": f"{date} 09:30:00",
|
|
"end_date": f"{date} 15:00:00",
|
|
},
|
|
"ts_code,trade_time,open,high,low,close,vol,amount",
|
|
unit="yuan/share",
|
|
)
|
|
|
|
def snapshot_inputs(
|
|
self, trade_date: str, previous_trade_date: str
|
|
) -> dict[str, ProviderResult | dict[str, Any]]:
|
|
current = _compact(trade_date)
|
|
previous = _compact(previous_trade_date)
|
|
daily = self._query(
|
|
"daily",
|
|
{"trade_date": current},
|
|
"ts_code,trade_date,open,high,low,close,pre_close,pct_chg,vol,amount",
|
|
unit="mixed",
|
|
)
|
|
event_fields = (
|
|
"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"
|
|
)
|
|
datasets: dict[str, ProviderResult | dict[str, Any]] = {"daily": daily}
|
|
datasets["price_limits"] = self._query(
|
|
"stk_limit",
|
|
{"trade_date": current},
|
|
"ts_code,trade_date,up_limit,down_limit",
|
|
unit="yuan/share",
|
|
)
|
|
for key, limit_type, date in (
|
|
("limit_up", "U", current),
|
|
("limit_down", "D", current),
|
|
("broken", "Z", current),
|
|
("previous_limit_up", "U", previous),
|
|
):
|
|
datasets[key] = self._query(
|
|
"limit_list_d",
|
|
{"trade_date": date, "limit_type": limit_type},
|
|
event_fields,
|
|
unit="mixed",
|
|
empty_is_complete=True,
|
|
)
|
|
return datasets
|
|
|
|
def sector_members(self, representative: str, trade_date: str) -> ProviderResult:
|
|
target = _compact(trade_date)
|
|
memberships = self._membership_rows({"ts_code": representative})
|
|
active = [row for row in memberships if _active_on(row, target)]
|
|
if not active:
|
|
raise ProviderError("未找到该股票在目标日期的申万行业")
|
|
industry = max(
|
|
active,
|
|
key=lambda row: (
|
|
str(row.get("in_date") or ""),
|
|
str(row.get("l2_code") or ""),
|
|
),
|
|
)
|
|
sector_code = str(industry.get("l2_code") or "")
|
|
sector_name = str(industry.get("l2_name") or "").strip()
|
|
if not sector_code:
|
|
raise ProviderError("该股票缺少申万二级行业")
|
|
members = [
|
|
row
|
|
for row in self._membership_rows({"l2_code": sector_code})
|
|
if _active_on(row, target)
|
|
]
|
|
deduplicated: dict[str, dict[str, Any]] = {}
|
|
for row in members:
|
|
code = str(row.get("ts_code") or "")
|
|
current = deduplicated.get(code)
|
|
if code and (
|
|
current is None
|
|
or str(row.get("in_date") or "") > str(current.get("in_date") or "")
|
|
):
|
|
deduplicated[code] = row
|
|
if not deduplicated:
|
|
raise ProviderError("该申万行业没有有效成分股")
|
|
daily = self._query(
|
|
"daily",
|
|
{"trade_date": target},
|
|
"ts_code,trade_date,open,close,pct_chg,amount",
|
|
unit="mixed",
|
|
)
|
|
quote_map = {str(row.get("ts_code") or ""): row for row in daily.rows}
|
|
rows = []
|
|
for code, member in deduplicated.items():
|
|
quote = quote_map.get(code) or {}
|
|
rows.append(
|
|
{
|
|
"sector_code": sector_code,
|
|
"sector_name": sector_name,
|
|
"ts_code": code,
|
|
"name": str(member.get("name") or "").strip(),
|
|
"change": _number(quote.get("pct_chg")) if quote else None,
|
|
"open": _number(quote.get("open")) if quote else None,
|
|
"close": _number(quote.get("close")) if quote else None,
|
|
"amount": _number(quote.get("amount")) * 1000 if quote else None,
|
|
"quoted": bool(quote),
|
|
}
|
|
)
|
|
coverage = sum(bool(row["quoted"]) for row in rows) / len(rows)
|
|
return ProviderResult(tuple(rows), _metadata(self.source, "mixed", coverage))
|
|
|
|
def _membership_rows(self, params: dict[str, str]) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
fields = (
|
|
"l1_code,l1_name,l2_code,l2_name,l3_code,l3_name,"
|
|
"ts_code,name,in_date,out_date,is_new"
|
|
)
|
|
for is_new in ("Y", "N"):
|
|
result = self._query(
|
|
"index_member_all",
|
|
{**params, "is_new": is_new},
|
|
fields,
|
|
unit="membership",
|
|
empty_is_complete=True,
|
|
)
|
|
rows.extend(result.rows)
|
|
return rows
|
|
|
|
def _query(
|
|
self,
|
|
api_name: str,
|
|
params: dict[str, Any],
|
|
fields: str,
|
|
*,
|
|
unit: str,
|
|
adjustment: str = "not_applicable",
|
|
empty_is_complete: bool = False,
|
|
) -> ProviderResult:
|
|
if not self.configured:
|
|
raise ProviderError("行情服务尚未配置")
|
|
token = self._token()
|
|
if not token:
|
|
raise ProviderError("行情服务尚未配置")
|
|
body = json.dumps(
|
|
{"api_name": api_name, "token": token, "params": params, "fields": fields},
|
|
ensure_ascii=False,
|
|
).encode("utf-8")
|
|
request = urllib.request.Request(
|
|
self.url,
|
|
data=body,
|
|
headers={"Content-Type": "application/json", "User-Agent": "XiaobaiReview/2"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self._timeout) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
except (urllib.error.URLError, TimeoutError, OSError, json.JSONDecodeError) as exc:
|
|
raise ProviderError("行情服务请求失败") from exc
|
|
if payload.get("code") not in (None, 0):
|
|
raise ProviderError(str(payload.get("msg") or "行情服务拒绝请求"))
|
|
data = payload.get("data") or {}
|
|
columns = data.get("fields") or []
|
|
rows = tuple(dict(zip(columns, item, strict=False)) for item in data.get("items") or [])
|
|
coverage = 1 if rows or empty_is_complete else 0
|
|
return ProviderResult(rows, _metadata(self.source, unit, coverage, adjustment))
|
|
|
|
def _token(self) -> str:
|
|
return str(self._token_provider() or "").strip()
|
|
|
|
|
|
def _metadata(
|
|
source: DataSource, unit: str, coverage: float, adjustment: str = "not_applicable"
|
|
) -> ObservationMetadata:
|
|
return ObservationMetadata(
|
|
source=source,
|
|
observed_at=datetime.now(SHANGHAI),
|
|
unit=unit,
|
|
adjustment=adjustment,
|
|
freshness_seconds=0,
|
|
coverage=coverage,
|
|
state=SnapshotState.ARCHIVE,
|
|
usage=DataUsage.CALCULATION,
|
|
)
|
|
|
|
|
|
def _compact(value: str) -> str:
|
|
normalized = value.replace("-", "")
|
|
if len(normalized) != 8 or not normalized.isdigit():
|
|
raise ProviderError("日期格式无效")
|
|
return normalized
|
|
|
|
|
|
def _display(value: str) -> str:
|
|
compact = _compact(value)
|
|
return f"{compact[:4]}-{compact[4:6]}-{compact[6:]}"
|
|
|
|
|
|
def _active_on(row: dict[str, Any], trade_date: str) -> bool:
|
|
start = str(row.get("in_date") or "")
|
|
end = str(row.get("out_date") or "")
|
|
return (not start or start <= trade_date) and (not end or end > trade_date)
|
|
|
|
|
|
def _number(value: Any) -> float:
|
|
try:
|
|
number = float(value)
|
|
return number if number == number else 0.0
|
|
except (TypeError, ValueError):
|
|
return 0.0
|