polynode 0.5.5__py3-none-any.whl

polynode/engine.py

@@ -0,0 +1,201 @@
+"""High-level orderbook engine with local state and filtered views."""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from .orderbook import OrderbookWS
+from .orderbook_state import LocalOrderbook
+from .types.orderbook import (
+    BookSnapshot,
+    BookUpdate,
+    OrderbookLevel,
+    PriceChange,
+)
+
+
+class EngineView:
+    """A filtered view over the shared orderbook state for specific tokens."""
+
+    def __init__(
+        self,
+        token_ids: list[str],
+        state: LocalOrderbook,
+        remove_from_engine: Callable,
+    ) -> None:
+        self._tokens: set[str] = set(token_ids)
+        self._state = state
+        self._remove_from_engine = remove_from_engine
+        self._destroyed = False
+        self._handlers: dict[str, list[Callable]] = {"update": [], "price": []}
+
+    def on(self, event: str, handler: Callable) -> EngineView:
+        if not self._destroyed and event in self._handlers:
+            self._handlers[event].append(handler)
+        return self
+
+    def off(self, event: str, handler: Callable) -> EngineView:
+        if event in self._handlers:
+            try:
+                self._handlers[event].remove(handler)
+            except ValueError:
+                pass
+        return self
+
+    def set_tokens(self, token_ids: list[str]) -> None:
+        self._tokens = set(token_ids)
+
+    def book(self, token_id: str) -> tuple[list[OrderbookLevel], list[OrderbookLevel]] | None:
+        return self._state.get_book(token_id)
+
+    def midpoint(self, token_id: str) -> float | None:
+        return _compute_midpoint(self._state, token_id)
+
+    def spread(self, token_id: str) -> float | None:
+        return self._state.get_spread(token_id)
+
+    def best_bid(self, token_id: str) -> OrderbookLevel | None:
+        return self._state.get_best_bid(token_id)
+
+    def best_ask(self, token_id: str) -> OrderbookLevel | None:
+        return self._state.get_best_ask(token_id)
+
+    def destroy(self) -> None:
+        self._destroyed = True
+        self._handlers = {"update": [], "price": []}
+        self._tokens.clear()
+        self._remove_from_engine()
+
+    def _emit_update(self, data: BookSnapshot | BookUpdate) -> None:
+        if self._destroyed:
+            return
+        for h in self._handlers["update"]:
+            h(data)
+
+    def _emit_price(self, data: PriceChange) -> None:
+        if self._destroyed:
+            return
+        for h in self._handlers["price"]:
+            h(data)
+
+
+class OrderbookEngine:
+    """Manages one WebSocket connection, local state, and multiple filtered views."""
+
+    def __init__(
+        self,
+        api_key: str,
+        *,
+        ob_url: str = "wss://ob.polynode.dev/ws",
+        compress: bool = True,
+        auto_reconnect: bool = True,
+        max_reconnect_attempts: int = 0,
+        reconnect_base_delay: float = 1.0,
+        reconnect_max_delay: float = 30.0,
+    ) -> None:
+        self._ws = OrderbookWS(
+            api_key,
+            ob_url,
+            compress=compress,
+            auto_reconnect=auto_reconnect,
+            max_reconnect_attempts=max_reconnect_attempts,
+            reconnect_base_delay=reconnect_base_delay,
+            reconnect_max_delay=reconnect_max_delay,
+        )
+        self._state = LocalOrderbook()
+        self._views: set[EngineView] = set()
+        self._handlers: dict[str, list[Callable]] = {"update": [], "price": [], "ready": []}
+
+        self._ws.on("snapshot", self._on_snapshot)
+        self._ws.on("update", self._on_update)
+        self._ws.on("price", self._on_price)
+        self._ws.on("snapshots_done", self._on_ready)
+
+    def _on_snapshot(self, snap: BookSnapshot) -> None:
+        self._state.apply_snapshot(snap)
+        self._route_update(snap)
+
+    def _on_update(self, update: BookUpdate) -> None:
+        self._state.apply_update(update)
+        self._route_update(update)
+
+    def _on_price(self, change: PriceChange) -> None:
+        for h in self._handlers["price"]:
+            h(change)
+        for view in self._views:
+            if view._destroyed:
+                continue
+            if any(a.asset_id in view._tokens for a in change.assets):
+                view._emit_price(change)
+
+    def _on_ready(self, msg: Any) -> None:
+        for h in self._handlers["ready"]:
+            h()
+
+    async def subscribe(self, identifiers: list[str]) -> None:
+        await self._ws.subscribe(identifiers)
+
+    def on(self, event: str, handler: Callable) -> OrderbookEngine:
+        if event in self._handlers:
+            self._handlers[event].append(handler)
+        return self
+
+    def off(self, event: str, handler: Callable) -> OrderbookEngine:
+        if event in self._handlers:
+            try:
+                self._handlers[event].remove(handler)
+            except ValueError:
+                pass
+        return self
+
+    def view(self, token_ids: list[str]) -> EngineView:
+        v = EngineView(token_ids, self._state, lambda: self._views.discard(v))
+        self._views.add(v)
+        return v
+
+    def book(self, token_id: str) -> tuple[list[OrderbookLevel], list[OrderbookLevel]] | None:
+        return self._state.get_book(token_id)
+
+    def midpoint(self, token_id: str) -> float | None:
+        return _compute_midpoint(self._state, token_id)
+
+    def spread(self, token_id: str) -> float | None:
+        return self._state.get_spread(token_id)
+
+    def best_bid(self, token_id: str) -> OrderbookLevel | None:
+        return self._state.get_best_bid(token_id)
+
+    def best_ask(self, token_id: str) -> OrderbookLevel | None:
+        return self._state.get_best_ask(token_id)
+
+    @property
+    def connection(self) -> OrderbookWS:
+        return self._ws
+
+    @property
+    def size(self) -> int:
+        return self._state.size
+
+    def close(self) -> None:
+        for view in list(self._views):
+            view.destroy()
+        self._views.clear()
+        self._state.clear()
+        self._ws.disconnect()
+
+    def _route_update(self, update: BookSnapshot | BookUpdate) -> None:
+        for h in self._handlers["update"]:
+            h(update)
+        for view in self._views:
+            if view._destroyed:
+                continue
+            if update.asset_id in view._tokens:
+                view._emit_update(update)
+
+
+def _compute_midpoint(state: LocalOrderbook, token_id: str) -> float | None:
+    bid = state.get_best_bid(token_id)
+    ask = state.get_best_ask(token_id)
+    if not bid or not ask:
+        return None
+    return (float(bid.price) + float(ask.price)) / 2

polynode/errors.py

@@ -0,0 +1,35 @@
+"""PolyNode SDK error hierarchy."""
+
+from __future__ import annotations
+
+
+class PolyNodeError(Exception):
+    """Base error for all PolyNode SDK errors."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(message)
+
+
+class ApiError(PolyNodeError):
+    """HTTP API error with status code."""
+
+    def __init__(self, message: str, status: int) -> None:
+        self.status = status
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        return f"ApiError({self.status}): {self.message}"
+
+
+class WsError(PolyNodeError):
+    """WebSocket error with optional error code."""
+
+    def __init__(self, message: str, code: str | None = None) -> None:
+        self.code = code
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        if self.code:
+            return f"WsError({self.code}): {self.message}"
+        return f"WsError: {self.message}"

polynode/orderbook.py

@@ -0,0 +1,243 @@
+"""Async WebSocket client for real-time orderbook updates."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import math
+import random
+import zlib
+from typing import Any, Callable
+
+import websockets
+import websockets.asyncio.client
+
+from .errors import WsError
+from .types.orderbook import (
+    BookSnapshot,
+    BookUpdate,
+    ObErrorMessage,
+    PriceChange,
+    SnapshotsDoneMessage,
+    SubscribedMessage,
+)
+
+
+class OrderbookWS:
+    """Async WebSocket client for orderbook streaming from ob.polynode.dev."""
+
+    def __init__(self, api_key: str, ob_url: str, *, compress: bool = True, auto_reconnect: bool = True,
+                 max_reconnect_attempts: int = 0, reconnect_base_delay: float = 1.0,
+                 reconnect_max_delay: float = 30.0) -> None:
+        self._api_key = api_key
+        self._ob_url = ob_url
+        self._compress = compress
+        self._auto_reconnect = auto_reconnect
+        self._max_reconnect = max_reconnect_attempts
+        self._base_delay = reconnect_base_delay
+        self._max_delay = reconnect_max_delay
+
+        self._socket: Any = None
+        self._connected = False
+        self._intentional_close = False
+        self._reconnect_attempts = 0
+        self._token_ids: list[str] = []
+        self._recv_task: asyncio.Task | None = None
+
+        self._handlers: dict[str, list[Callable]] = {
+            "snapshot": [],
+            "update": [],
+            "price": [],
+            "snapshots_done": [],
+            "subscribed": [],
+            "error": [],
+            "*": [],
+        }
+        self._on_connect: list[Callable] = []
+        self._on_disconnect: list[Callable] = []
+        self._on_reconnect: list[Callable] = []
+        self._on_error: list[Callable] = []
+
+    @property
+    def is_connected(self) -> bool:
+        return self._connected
+
+    def on_connect(self, handler: Callable) -> OrderbookWS:
+        self._on_connect.append(handler)
+        return self
+
+    def on_disconnect(self, handler: Callable) -> OrderbookWS:
+        self._on_disconnect.append(handler)
+        return self
+
+    def on_reconnect(self, handler: Callable) -> OrderbookWS:
+        self._on_reconnect.append(handler)
+        return self
+
+    def on_error(self, handler: Callable) -> OrderbookWS:
+        self._on_error.append(handler)
+        return self
+
+    def on(self, event: str, handler: Callable) -> OrderbookWS:
+        if event in self._handlers:
+            self._handlers[event].append(handler)
+        return self
+
+    def off(self, event: str, handler: Callable) -> OrderbookWS:
+        if event in self._handlers:
+            try:
+                self._handlers[event].remove(handler)
+            except ValueError:
+                pass
+        return self
+
+    async def subscribe(self, token_ids: list[str]) -> None:
+        self._token_ids = token_ids
+        await self._ensure_connected()
+        await self._send(json.dumps({"action": "subscribe", "markets": token_ids}))
+
+    def unsubscribe(self) -> None:
+        self._token_ids = []
+        if self._connected and self._socket:
+            asyncio.ensure_future(self._send(json.dumps({"action": "unsubscribe"})))
+
+    async def connect(self) -> None:
+        if self._connected and self._socket:
+            return
+        await self._connect()
+
+    def disconnect(self) -> None:
+        self._intentional_close = True
+        if self._recv_task and not self._recv_task.done():
+            self._recv_task.cancel()
+        if self._socket:
+            asyncio.ensure_future(self._socket.close())
+            self._socket = None
+        self._connected = False
+
+    # ── Private ──
+
+    async def _ensure_connected(self) -> None:
+        if not self._connected or not self._socket:
+            await self._connect()
+
+    async def _connect(self) -> None:
+        url = f"{self._ob_url}?key={self._api_key}"
+        if self._compress:
+            url += "&compress=zlib"
+
+        self._socket = await websockets.asyncio.client.connect(url, max_size=10 * 1024 * 1024)
+        self._connected = True
+        self._reconnect_attempts = 0
+        self._intentional_close = False
+
+        for h in self._on_connect:
+            h()
+
+        self._recv_task = asyncio.ensure_future(self._recv_loop())
+
+    async def _recv_loop(self) -> None:
+        try:
+            async for raw in self._socket:
+                try:
+                    if isinstance(raw, bytes) and self._compress:
+                        text = zlib.decompress(raw, -zlib.MAX_WBITS).decode("utf-8")
+                    elif isinstance(raw, bytes):
+                        text = raw.decode("utf-8")
+                    else:
+                        text = raw
+                    self._handle_message(text)
+                except Exception as e:
+                    err = WsError(f"Failed to process message: {e}")
+                    for h in self._on_error:
+                        h(err)
+        except websockets.exceptions.ConnectionClosed:
+            pass
+        except asyncio.CancelledError:
+            return
+        finally:
+            self._connected = False
+            reason = "closed" if self._intentional_close else "lost"
+            for h in self._on_disconnect:
+                h(reason)
+            if not self._intentional_close and self._auto_reconnect:
+                asyncio.ensure_future(self._schedule_reconnect())
+
+    async def _send(self, data: str) -> None:
+        if self._socket and self._connected:
+            await self._socket.send(data)
+
+    def _handle_message(self, text: str) -> None:
+        msg = json.loads(text)
+
+        if msg.get("error"):
+            err_msg = ObErrorMessage.model_validate(msg)
+            for h in self._handlers["error"]:
+                h(err_msg)
+            return
+
+        msg_type = msg.get("type")
+
+        if msg_type == "subscribed":
+            sub_msg = SubscribedMessage(markets=msg.get("markets", 0))
+            for h in self._handlers["subscribed"]:
+                h(sub_msg)
+
+        elif msg_type == "snapshot_batch":
+            snapshots = msg.get("snapshots", [])
+            for snap_data in snapshots:
+                snap = BookSnapshot.model_validate(snap_data)
+                for h in self._handlers["snapshot"]:
+                    h(snap)
+                for h in self._handlers["*"]:
+                    h(snap)
+
+        elif msg_type == "snapshots_done":
+            done = SnapshotsDoneMessage(total=msg.get("total", 0))
+            for h in self._handlers["snapshots_done"]:
+                h(done)
+
+        elif msg_type == "batch":
+            updates = msg.get("updates", [])
+            for update_data in updates:
+                self._dispatch_update(update_data)
+
+    def _dispatch_update(self, data: dict) -> None:
+        update_type = data.get("type")
+        if update_type == "book_snapshot":
+            snap = BookSnapshot.model_validate(data)
+            for h in self._handlers["snapshot"]:
+                h(snap)
+            for h in self._handlers["*"]:
+                h(snap)
+        elif update_type == "book_update":
+            update = BookUpdate.model_validate(data)
+            for h in self._handlers["update"]:
+                h(update)
+            for h in self._handlers["*"]:
+                h(update)
+        elif update_type == "price_change":
+            change = PriceChange.model_validate(data)
+            for h in self._handlers["price"]:
+                h(change)
+            for h in self._handlers["*"]:
+                h(change)
+
+    async def _schedule_reconnect(self) -> None:
+        if self._max_reconnect and self._reconnect_attempts >= self._max_reconnect:
+            return
+
+        self._reconnect_attempts += 1
+        delay = min(self._base_delay * math.pow(2, self._reconnect_attempts - 1), self._max_delay)
+        jitter = delay * (0.5 + random.random() * 0.5)
+
+        await asyncio.sleep(jitter)
+
+        try:
+            await self._connect()
+            for h in self._on_reconnect:
+                h(self._reconnect_attempts)
+            if self._token_ids:
+                await self._send(json.dumps({"action": "subscribe", "markets": self._token_ids}))
+        except Exception:
+            pass

polynode/orderbook_state.py

@@ -0,0 +1,77 @@
+"""Local orderbook state manager."""
+
+from __future__ import annotations
+
+from .types.orderbook import BookSnapshot, BookUpdate, OrderbookLevel, PriceChange
+
+
+class LocalOrderbook:
+    """Maintains sorted local copies of orderbooks for any number of assets."""
+
+    def __init__(self) -> None:
+        self._books: dict[str, tuple[list[OrderbookLevel], list[OrderbookLevel]]] = {}
+
+    def apply_snapshot(self, snap: BookSnapshot) -> None:
+        bids = sorted(list(snap.bids), key=lambda l: float(l.price), reverse=True)
+        asks = sorted(list(snap.asks), key=lambda l: float(l.price))
+        self._books[snap.asset_id] = (bids, asks)
+
+    def apply_update(self, update: BookUpdate) -> None:
+        if update.asset_id not in self._books:
+            self._books[update.asset_id] = ([], [])
+        bids, asks = self._books[update.asset_id]
+        new_bids = _apply_deltas(bids, update.bids, reverse=True)
+        new_asks = _apply_deltas(asks, update.asks, reverse=False)
+        self._books[update.asset_id] = (new_bids, new_asks)
+
+    def apply_price_change(self, change: PriceChange) -> None:
+        pass
+
+    def get_book(self, asset_id: str) -> tuple[list[OrderbookLevel], list[OrderbookLevel]] | None:
+        return self._books.get(asset_id)
+
+    def get_best_bid(self, asset_id: str) -> OrderbookLevel | None:
+        book = self._books.get(asset_id)
+        if book and book[0]:
+            return book[0][0]
+        return None
+
+    def get_best_ask(self, asset_id: str) -> OrderbookLevel | None:
+        book = self._books.get(asset_id)
+        if book and book[1]:
+            return book[1][0]
+        return None
+
+    def get_spread(self, asset_id: str) -> float | None:
+        bid = self.get_best_bid(asset_id)
+        ask = self.get_best_ask(asset_id)
+        if not bid or not ask:
+            return None
+        return float(ask.price) - float(bid.price)
+
+    @property
+    def size(self) -> int:
+        return len(self._books)
+
+    def clear(self) -> None:
+        self._books.clear()
+
+
+def _apply_deltas(
+    levels: list[OrderbookLevel],
+    deltas: list[OrderbookLevel],
+    *,
+    reverse: bool,
+) -> list[OrderbookLevel]:
+    price_map: dict[str, str] = {}
+    for level in levels:
+        price_map[level.price] = level.size
+    for delta in deltas:
+        if delta.size == "0" or delta.size == "0.00":
+            price_map.pop(delta.price, None)
+        else:
+            price_map[delta.price] = delta.size
+
+    result = [OrderbookLevel(price=p, size=s) for p, s in price_map.items()]
+    result.sort(key=lambda l: float(l.price), reverse=reverse)
+    return result