"""底仓、补仓记录与待确认订单的 JSON 存储。""" from __future__ import annotations import json import logging as log import math from dataclasses import asdict, dataclass, replace from pathlib import Path from threading import RLock from typing import Iterable from sdk import OrderItem, PositionItem TERMINAL_STATUSES = {"53", "54", "56", "57"} @dataclass(frozen=True, slots=True) class StateItem: code: str base_order_id: str = "" base_qty: int = 0 base_cost: float = 0.0 added_order_id: str = "" added_num: int = 0 added_qty: int = 0 added_cost: float = 0.0 @dataclass(frozen=True, slots=True) class PendingOrder: order_id: str code: str kind: str pre_qty: int class State: def __init__(self, path: str | Path) -> None: self.path = Path(path) self.lock = RLock() self.items: dict[str, StateItem] = {} self.pending: dict[str, PendingOrder] = {} self.IsModify = False self._load() def _load(self) -> None: """启动时读取 JSON 中的持仓记录和待确认订单。""" if self.path.is_file(): raw = json.loads(self.path.read_text(encoding="utf-8")) self.items = {code: StateItem(**item) for code, item in raw["items"].items()} self.pending = {code: PendingOrder(**item) for code, item in raw["pending"].items()} @classmethod def for_strategy(cls, data_dir: str | Path, strategy: str, account_id: str) -> State: return cls(Path(data_dir) / f"{strategy}_{account_id}_state.json") def get(self, code: str) -> StateItem: with self.lock: return self.items[code] def busy(self, code: str) -> bool: with self.lock: return code in self.pending def new_order(self, order: PendingOrder) -> bool: """提交前保存待确认订单,同一证券已有 pending 时跳过。""" with self.lock: if order.code in self.pending: return False was_modified = self.IsModify self.pending[order.code] = order self.IsModify = True try: self.save() except Exception: del self.pending[order.code] self.IsModify = was_modified raise return True def reconcile(self, positions: Iterable[PositionItem], orders: list[OrderItem]) -> None: """核对 pending 的订单结果,再导入/核对实际持仓。""" by_id: dict[str, dict[str, OrderItem]] = {} for order in orders: if order.local_order_id: by_id.setdefault(order.local_order_id, {})[order.id] = order with self.lock: for code, pending in list(self.pending.items()): rows = [order for order in by_id.get(pending.order_id, {}).values() if order.code == code and order.side == "BUY"] if not rows: del self.pending[code] self.IsModify = True log.info("[状态] 清理无对应订单的 pending,代码=%s,订单=%s", code, pending.order_id) continue filled = _finished_fill(rows, pending.pre_qty) if filled is None: continue qty, amount = filled if qty: item = self.items.get(code, StateItem(code)) if pending.kind == "base": item = replace(item, base_order_id=pending.order_id, base_qty=qty, base_cost=amount / qty) else: total_qty = item.added_qty + qty total_amount = item.added_qty * item.added_cost + amount item = replace(item, added_order_id=pending.order_id, added_num=item.added_num + 1, added_qty=total_qty, added_cost=total_amount / total_qty) self.items[code] = item del self.pending[code] self.IsModify = True log.info("[状态] 对账结束,代码=%s,订单=%s,计划=%d,成交=%d", code, pending.order_id, pending.pre_qty, qty) for position in positions: code = position.stock_code if not code or position.volume <= 0: continue if code not in self.items and code not in self.pending: self.items[code] = StateItem(code=code, base_order_id=position.trade_id, base_qty=position.volume, base_cost=position.open_price) self.IsModify = True elif code in self.items: item = self.items[code] if item.base_qty + item.added_qty != position.volume: log.info("[状态] 持仓数量差异,代码=%s,记录=%d,实际=%d", code, item.base_qty + item.added_qty, position.volume) self.save() def save(self) -> None: with self.lock: if not self.IsModify: return self.path.parent.mkdir(parents=True, exist_ok=True) temporary = self.path.with_suffix(self.path.suffix + ".tmp") payload = dict(items={code: asdict(item) for code, item in self.items.items()}, pending={code: asdict(item) for code, item in self.pending.items()}) temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8") temporary.replace(self.path) self.IsModify = False def _finished_fill(orders: list[OrderItem], pre_qty: int) -> tuple[int, float] | None: """订单全部结束且成交金额完整后一次记账,部分成交撤单也按实计入。""" if sum(order.volume for order in orders) != pre_qty: return None qty, amount = 0, 0.0 for order in orders: if (order.status not in TERMINAL_STATUSES or (order.status == "56" and order.traded_volume != order.volume)): return None if order.traded_volume: value = order.trade_amount if order.trade_amount > 0 else order.trade_price * order.traded_volume if not math.isfinite(value) or value <= 0: return None qty += order.traded_volume amount += value return qty, amount