1"""Single-asset limit order book with price-time priority.23The book only knows orders and prices. Fees, accounts, latency and halts live in `engine.py`.4"""5from __future__ import annotations67from collections import deque8from dataclasses import dataclass9from typing import Iterator1011from sortedcontainers import SortedDict1213from .orders import Order, OrderStatus, OrderType, Side141516@dataclass(slots=True)17class Fill:18 resting: Order19 incoming: Order20 price: int21 qty: int222324class LimitOrderBook:25 __slots__ = ("asset", "bids", "asks", "_index", "_depth_cache")2627 def __init__(self, asset: int) -> None:28 self.asset = asset29 # bids: highest price first -> store negative price keys so .peekitem(0) is the best bid30 self.bids: SortedDict[int, deque[Order]] = SortedDict()31 self.asks: SortedDict[int, deque[Order]] = SortedDict()32 self._index: dict[int, Order] = {}3334 # ---- queries -------------------------------------------------------------------------35 def best_bid(self) -> int | None:36 return -self.bids.peekitem(0)[0] if self.bids else None3738 def best_ask(self) -> int | None:39 return self.asks.peekitem(0)[0] if self.asks else None4041 def mid(self) -> float | None:42 b, a = self.best_bid(), self.best_ask()43 if b is None or a is None:44 return None45 return (a + b) / 2.04647 def spread(self) -> int | None:48 b, a = self.best_bid(), self.best_ask()49 if b is None or a is None:50 return None51 return a - b5253 def depth(self, levels: int = 5) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:54 """Aggregated (price, qty) for the best `levels` on each side."""55 bids = []56 for i, (negp, q) in enumerate(self.bids.items()):57 if i >= levels:58 break59 bids.append((-negp, sum(o.remaining for o in q)))60 asks = []61 for i, (p, q) in enumerate(self.asks.items()):62 if i >= levels:63 break64 asks.append((p, sum(o.remaining for o in q)))65 return bids, asks6667 def depth_within(self, side: Side, ticks: int) -> int:68 """Total resting quantity within `ticks` of the best price on `side`."""69 if side is Side.BUY:70 bb = self.best_bid()71 if bb is None:72 return 073 return sum(o.remaining for negp, q in self.bids.items() if -negp >= bb - ticks for o in q)74 ba = self.best_ask()75 if ba is None:76 return 077 return sum(o.remaining for p, q in self.asks.items() if p <= ba + ticks for o in q)7879 def get(self, order_id: int) -> Order | None:80 return self._index.get(order_id)8182 def __contains__(self, order_id: int) -> bool:83 return order_id in self._index8485 def __len__(self) -> int:86 return len(self._index)8788 def resting_orders(self) -> Iterator[Order]:89 return iter(self._index.values())9091 def resting_for(self, agent_id: int) -> list[Order]:92 return [o for o in self._index.values() if o.agent_id == agent_id]9394 def total_qty(self) -> tuple[int, int]:95 return (sum(o.remaining for q in self.bids.values() for o in q),96 sum(o.remaining for q in self.asks.values() for o in q))9798 # ---- mutation ------------------------------------------------------------------------99 def _rest(self, order: Order) -> None:100 assert order.order_type is OrderType.LIMIT and order.price is not None101 book = self.bids if order.side is Side.BUY else self.asks102 key = -order.price if order.side is Side.BUY else order.price103 q = book.get(key)104 if q is None:105 q = deque()106 book[key] = q107 q.append(order)108 self._index[order.order_id] = order109 order.status = OrderStatus.OPEN110111 def _remove_resting(self, order: Order) -> None:112 book = self.bids if order.side is Side.BUY else self.asks113 key = -order.price if order.side is Side.BUY else order.price114 q = book[key]115 q.remove(order)116 if not q:117 del book[key]118 del self._index[order.order_id]119120 def match(self, incoming: Order, max_qty: int | None = None) -> list[Fill]:121 """Walk the opposite side while the incoming order crosses.122123 Returns the fills in execution order; mutates `incoming.remaining` and the resting orders.124 Does *not* rest the remainder (the engine decides, because it must reserve cash first).125 `max_qty` caps the executable quantity (used for market buys with a cash constraint).126 """127 fills: list[Fill] = []128 opp = self.asks if incoming.side is Side.BUY else self.bids129 budget = incoming.remaining if max_qty is None else min(incoming.remaining, max_qty)130 while budget > 0 and opp:131 key, queue = opp.peekitem(0)132 best_price = key if incoming.side is Side.BUY else -key133 if incoming.order_type is OrderType.LIMIT:134 assert incoming.price is not None135 if incoming.side is Side.BUY and best_price > incoming.price:136 break137 if incoming.side is Side.SELL and best_price < incoming.price:138 break139 while budget > 0 and queue:140 resting = queue[0]141 q = min(resting.remaining, budget)142 resting.remaining -= q143 incoming.remaining -= q144 budget -= q145 fills.append(Fill(resting=resting, incoming=incoming, price=best_price, qty=q))146 if resting.remaining == 0:147 resting.status = OrderStatus.FILLED148 queue.popleft()149 del self._index[resting.order_id]150 if not queue:151 del opp[key]152 if incoming.remaining == 0:153 incoming.status = OrderStatus.FILLED154 return fills155156 def rest(self, order: Order) -> None:157 self._rest(order)158159 def cancel(self, order_id: int) -> Order | None:160 order = self._index.get(order_id)161 if order is None:162 return None163 self._remove_resting(order)164 order.status = OrderStatus.CANCELLED165 return order166167 def clear_all(self) -> list[Order]:168 orders = list(self._index.values())169 for o in orders:170 self._remove_resting(o)171 o.status = OrderStatus.CANCELLED172 return orders173174 # ---- invariants (used by tests) ------------------------------------------------------175 def check_invariants(self) -> None:176 bb, ba = self.best_bid(), self.best_ask()177 if bb is not None and ba is not None:178 assert bb < ba, f"crossed book: bid {bb} >= ask {ba}"179 n = 0180 for negp, q in self.bids.items():181 assert q, "empty price level"182 for o in q:183 assert o.side is Side.BUY and o.price == -negp and o.remaining > 0 and o.status is OrderStatus.OPEN184 n += 1185 for p, q in self.asks.items():186 assert q, "empty price level"187 for o in q:188 assert o.side is Side.SELL and o.price == p and o.remaining > 0 and o.status is OrderStatus.OPEN189 n += 1190 assert n == len(self._index), "index/book mismatch"191