big-qmt/py-client/sdk/client.py

import json
from dataclasses import asdict, is_dataclass
from typing import Any

import httpx

from .context import ContextMixin
from .data import DataMixin
from .errors import APIError
from .get import GetMixin
from .portfolio import PortfolioMixin
from .sys import SysMixin
from .trade import TradeMixin


class HTTPClient:
    """复用连接池的同步 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) -> HTTPClient:
        return self

    def __exit__(self, *_args: object) -> None:
        self.close()

    def set_account_type(self, account_type: str) -> HTTPClient:
        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


class Client(
    ContextMixin,
    GetMixin,
    PortfolioMixin,
    DataMixin,
    TradeMixin,
    SysMixin,
    HTTPClient,
):
    """Client for the API exposed by ``qmt_rest_new.py``."""

def _is_idempotent(method: str, path: str) -> bool:
    if method == "GET":
        return True
    return path == "/api/data/full_tick"