119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, is_dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from .errors import APIError, BusinessError
|
|
|
|
|
|
def csv_join(items: list[str]) -> str:
|
|
return ",".join(item.strip() for item in items if item.strip())
|
|
|
|
|
|
class Client:
|
|
"""复用连接池的同步 QMT HTTP 客户端。"""
|
|
|
|
def __init__(self, base_url: str, token: str, timeout: float = 15.0) -> None:
|
|
self.base_url = base_url.rstrip("/")
|
|
self.token = token
|
|
self.timeout = timeout if timeout > 0 else 15.0
|
|
self.account_type = "STOCK"
|
|
self.http = httpx.Client(
|
|
base_url=self.base_url,
|
|
headers={"X-Token": token, "Accept": "application/json"},
|
|
timeout=httpx.Timeout(self.timeout),
|
|
limits=httpx.Limits(max_connections=20, max_keepalive_connections=10),
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self.http.close()
|
|
|
|
def __enter__(self) -> "Client":
|
|
return self
|
|
|
|
def __exit__(self, *_args: object) -> None:
|
|
self.close()
|
|
|
|
def set_account_type(self, account_type: str) -> "Client":
|
|
if account_type.strip():
|
|
self.account_type = account_type
|
|
return self
|
|
|
|
def _request_bytes(self, method: str, path: str, body: Any = None) -> bytes:
|
|
if is_dataclass(body):
|
|
body = asdict(body)
|
|
attempts = 2 if _is_idempotent(method, path) else 1
|
|
response: httpx.Response | None = None
|
|
for attempt in range(attempts):
|
|
try:
|
|
response = self.http.request(method, path, json=body)
|
|
break
|
|
except (httpx.ConnectError, httpx.ReadTimeout):
|
|
if attempt + 1 == attempts:
|
|
raise
|
|
assert response is not None
|
|
if response.status_code >= 400:
|
|
try:
|
|
message = response.text
|
|
except (ValueError, AttributeError):
|
|
message = response.text.strip()
|
|
raise APIError(response.status_code, str(message))
|
|
if not response.content:
|
|
return b""
|
|
|
|
return response.content
|
|
|
|
def _get_bytes(self, path: str) -> bytes:
|
|
return self._request_bytes("GET", path)
|
|
|
|
def _post_bytes(self, path: str, body: Any = None) -> bytes:
|
|
return self._request_bytes("POST", path, {} if body is None else body)
|
|
|
|
def _get_json(self, path: str) -> Any:
|
|
content = self._get_bytes(path)
|
|
return self._decode_json(path, content)
|
|
|
|
def _post_json(self, path: str, body: Any = None) -> Any:
|
|
content = self._post_bytes(path, body)
|
|
return self._decode_json(path, content)
|
|
|
|
@staticmethod
|
|
def _decode_json(path: str, content: bytes) -> Any:
|
|
if not content:
|
|
return None
|
|
try:
|
|
return json.loads(content)
|
|
except ValueError as exc:
|
|
raise ValueError(
|
|
f"invalid JSON from {path}: {content[:512]!r}"
|
|
) from exc
|
|
|
|
def _get_field(self, path: str, key: str) -> Any:
|
|
return self._get_json(path).get(key)
|
|
|
|
def _post_field(self, path: str, body: Any, key: str) -> Any:
|
|
result = self._post_json(path, body)
|
|
if isinstance(result, dict) and result.get("error"):
|
|
raise BusinessError(result["error"])
|
|
return result.get(key, result) if key and isinstance(result, dict) else result
|
|
|
|
|
|
def _is_idempotent(method: str, path: str) -> bool:
|
|
if method == "GET":
|
|
return True
|
|
prefixes = (
|
|
"/api/v2/",
|
|
"/api/holding",
|
|
"/api/money/",
|
|
"/api/context/",
|
|
"/api/check/",
|
|
"/api/data/",
|
|
"/api/trade/trade_detail_data",
|
|
"/api/order/deal",
|
|
)
|
|
unsafe = ("subscribe", "unsubscribe")
|
|
return path.startswith(prefixes) and not any(word in path for word in unsafe)
|