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

from __future__ import annotations

from typing import Any

from .models import Assets, DealItem, OrderItem, Portfolio, PositionItem


class PortfolioMixin:
    def org(self, datatype: str) -> list[dict[str, Any]]:
        """查询 account、order、deal 或 position，返回原始字段字典列表。"""
        datatype = str(datatype).strip().lower()
        if datatype not in {"account", "order", "deal", "position"}:
            raise ValueError(f"unsupported org datatype: {datatype}")
        response = self._get_json(f"/api/portfolio/org/{datatype}")
        return response["data"]

    def portfolio(self) -> Portfolio:
        data = self._get_json("/api/portfolio") or {}
        positions = {
            code: PositionItem(**value)
            for code, value in data.get("positions", {}).items()
        }
        return Portfolio(
            assets=Assets(**data.get("assets", {})),
            positions=positions,
            orders=[OrderItem(**row) for row in data.get("orders", [])],
        )

    def positions(self) -> tuple[list[str], list[PositionItem]]:
        data = self._get_json("/api/portfolio/positions") or {}
        positions = [
            PositionItem(**value)
            for value in data.values()
        ]
        return list(data), positions

    def assets(self) -> Assets:
        return Assets(**(self._get_json("/api/portfolio/assets") or {}))

    def orders(self) -> list[OrderItem]:
        data = self._get_json("/api/portfolio/order") or []
        return [OrderItem(**row) for row in data]

    def deals(self) -> list[DealItem]:
        data = self._get_json("/api/portfolio/deal") or []
        return [DealItem(**row) for row in data]

    def trade_detail_data(self, datatype: str) -> Any:
        datatype = str(datatype).strip().lower()
        handlers = {
            "account": self.assets,
            "position": lambda: self.positions()[1],
            "order": self.orders,
            "deal": self.deals,
        }
        handler = handlers.get(datatype)
        if handler is None:
            raise ValueError(f"unsupported trade detail datatype: {datatype}")
        result = handler()
        return [result] if datatype == "account" else result