70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
import structlog
|
|
|
|
logger = structlog.get_logger()
|
|
|
|
|
|
@dataclass
|
|
class BookLevel:
|
|
price: float
|
|
size: float
|
|
|
|
@classmethod
|
|
def from_list(cls, data: list) -> "BookLevel":
|
|
return cls(price=float(data[0]), size=float(data[1]))
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"price": self.price, "size": self.size}
|
|
|
|
|
|
@dataclass
|
|
class OrderBookTop5:
|
|
symbol: str
|
|
bids: list[BookLevel] = field(default_factory=list)
|
|
asks: list[BookLevel] = field(default_factory=list)
|
|
ts_ms: int = 0
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"symbol": self.symbol,
|
|
"bids": [b.to_dict() for b in self.bids],
|
|
"asks": [a.to_dict() for a in self.asks],
|
|
"ts_ms": self.ts_ms,
|
|
}
|
|
|
|
|
|
class BookStore:
|
|
def __init__(self) -> None:
|
|
self._books: dict[str, OrderBookTop5] = {}
|
|
self._log = logger.bind(component="book_store")
|
|
|
|
def update(self, raw: dict) -> Optional[OrderBookTop5]:
|
|
topic = raw.get("topic", "")
|
|
data = raw.get("data", {})
|
|
|
|
topic_suffix = topic.split(":")[-1] if ":" in topic else ""
|
|
symbol = topic_suffix.split(",")[0].strip() if topic_suffix else ""
|
|
asks_raw = data.get("asks", [])
|
|
bids_raw = data.get("bids", [])
|
|
|
|
if not symbol:
|
|
return None
|
|
|
|
ts_ms = int(data.get("time", time.time() * 1000))
|
|
|
|
bids = [BookLevel.from_list(b) for b in bids_raw[:1]]
|
|
asks = [BookLevel.from_list(a) for a in asks_raw[:1]]
|
|
|
|
book = OrderBookTop5(symbol=symbol, bids=bids, asks=asks, ts_ms=ts_ms)
|
|
self._books[symbol] = book
|
|
|
|
return book
|
|
|
|
def get(self, symbol: str) -> Optional[OrderBookTop5]:
|
|
return self._books.get(symbol)
|
|
|
|
def get_all(self) -> dict[str, OrderBookTop5]:
|
|
return self._books.copy() |