53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
"""
|
|
Minimal async client for Tushare Pro API
|
|
"""
|
|
from typing import Any, Dict, List, Optional
|
|
import httpx
|
|
|
|
TUSHARE_PRO_URL = "https://api.tushare.pro"
|
|
|
|
|
|
class TushareClient:
|
|
def __init__(self, token: str):
|
|
self.token = token
|
|
self._client = httpx.AsyncClient(timeout=30)
|
|
|
|
async def query(
|
|
self,
|
|
api_name: str,
|
|
params: Optional[Dict[str, Any]] = None,
|
|
fields: Optional[str] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
payload = {
|
|
"api_name": api_name,
|
|
"token": self.token,
|
|
"params": params or {},
|
|
}
|
|
# default larger page size if not provided
|
|
if "limit" not in payload["params"]:
|
|
payload["params"]["limit"] = 5000
|
|
if fields:
|
|
payload["fields"] = fields
|
|
resp = await self._client.post(TUSHARE_PRO_URL, json=payload)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
if data.get("code") != 0:
|
|
err = data.get("msg") or "Tushare error"
|
|
raise RuntimeError(f"{api_name}: {err}")
|
|
fields_def = data.get("data", {}).get("fields", [])
|
|
items = data.get("data", {}).get("items", [])
|
|
rows: List[Dict[str, Any]] = []
|
|
for it in items:
|
|
row = {fields_def[i]: it[i] for i in range(len(fields_def))}
|
|
rows.append(row)
|
|
return rows
|
|
|
|
async def aclose(self):
|
|
await self._client.aclose()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
await self.aclose()
|