from __future__ import annotations import codecs import queue import threading from typing import Any, Callable, Iterable, Iterator, Literal, overload from ..exceptions import ConcurrencyError from ..frames import OP_BINARY, OP_CONT, OP_TEXT, Frame from ..typing import Data from .utils import Deadline __all__ = ["Assembler"] UTF8Decoder = codecs.getincrementaldecoder("utf-8") class Assembler: """ Assemble messages from frames. :class:`Assembler` expects only data frames. The stream of frames must respect the protocol; if it doesn't, the behavior is undefined. Args: pause: Called when the buffer of frames goes above the high water mark; should pause reading from the network. resume: Called when the buffer of frames goes below the low water mark; should resume reading from the network. """ def __init__( self, high: int | None = None, low: int | None = None, pause: Callable[[], Any] = lambda: None, resume: Callable[[], Any] = lambda: None, ) -> None: # Serialize reads and writes -- except for reads via synchronization # primitives provided by the threading and queue modules. self.mutex = threading.Lock() # Queue of incoming frames. self.frames: queue.SimpleQueue[Frame | None] = queue.SimpleQueue() # We cannot put a hard limit on the size of the queue because a single # call to Protocol.data_received() could produce thousands of frames, # which must be buffered. Instead, we pause reading when the buffer goes # above the high limit and we resume when it goes under the low limit. if high is not None and low is None: low = high // 4 if high is None and low is not None: high = low * 4 if high is not None and low is not None: if low < 0: raise ValueError("low must be positive or equal to zero") if high < low: raise ValueError("high must be greater than or equal to low") self.high, self.low = high, low self.pause = pause self.resume = resume self.paused = False # This flag prevents concurrent calls to get() by user code. self.get_in_progress = False # This flag marks the end of the connection. self.closed = False def get_next_frame(self, timeout: float | None = None) -> Frame: # Helper to factor out the logic for getting the next frame from the # queue, while handling timeouts and reaching the end of the stream. if self.closed: try: frame = self.frames.get(block=False) except queue.Empty: raise EOFError("stream of frames ended") from None else: try: # Check for a frame that's already received if timeout <= 0. # SimpleQueue.get() doesn't support negative timeout values. if timeout is not None and timeout <= 0: frame = self.frames.get(block=False) else: frame = self.frames.get(block=True, timeout=timeout) except queue.Empty: raise TimeoutError(f"timed out in {timeout:.1f}s") from None if frame is None: raise EOFError("stream of frames ended") return frame def reset_queue(self, frames: Iterable[Frame]) -> None: # Helper to put frames back into the queue after they were fetched. # This happens only when the queue is empty. However, by the time # we acquire self.mutex, put() may have added items in the queue. # Therefore, we must handle the case where the queue is not empty. frame: Frame | None with self.mutex: queued = [] try: while True: queued.append(self.frames.get(block=False)) except queue.Empty: pass for frame in frames: self.frames.put(frame) # This loop runs only when a race condition occurs. for frame in queued: # pragma: no cover self.frames.put(frame) # This overload structure is required to avoid the error: # "parameter without a default follows parameter with a default" @overload def get(self, timeout: float | None, decode: Literal[True]) -> str: ... @overload def get(self, timeout: float | None, decode: Literal[False]) -> bytes: ... @overload def get(self, timeout: float | None = None, *, decode: Literal[True]) -> str: ... @overload def get(self, timeout: float | None = None, *, decode: Literal[False]) -> bytes: ... @overload def get(self, timeout: float | None = None, decode: bool | None = None) -> Data: ... def get(self, timeout: float | None = None, decode: bool | None = None) -> Data: """ Read the next message. :meth:`get` returns a single :class:`str` or :class:`bytes`. If the message is fragmented, :meth:`get` waits until the last frame is received, then it reassembles the message and returns it. To receive messages frame by frame, use :meth:`get_iter` instead. Args: timeout: If a timeout is provided and elapses before a complete message is received, :meth:`get` raises :exc:`TimeoutError`. decode: :obj:`False` disables UTF-8 decoding of text frames and returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of binary frames and returns :class:`str`. Raises: EOFError: If the stream of frames has ended. UnicodeDecodeError: If a text frame contains invalid UTF-8. ConcurrencyError: If two coroutines run :meth:`get` or :meth:`get_iter` concurrently. TimeoutError: If a timeout is provided and elapses before a complete message is received. """ with self.mutex: if self.get_in_progress: raise ConcurrencyError("get() or get_iter() is already running") self.get_in_progress = True # Locking with get_in_progress prevents concurrent execution # until get() fetches a complete message or times out. try: deadline = Deadline(timeout) # First frame frame = self.get_next_frame(deadline.timeout(raise_if_elapsed=False)) with self.mutex: self.maybe_resume() assert frame.opcode is OP_TEXT or frame.opcode is OP_BINARY if decode is None: decode = frame.opcode is OP_TEXT frames = [frame] # Following frames, for fragmented messages while not frame.fin: try: frame = self.get_next_frame( deadline.timeout(raise_if_elapsed=False) ) except TimeoutError: # Put frames already received back into the queue # so that future calls to get() can return them. self.reset_queue(frames) raise with self.mutex: self.maybe_resume() assert frame.opcode is OP_CONT frames.append(frame) finally: self.get_in_progress = False data = b"".join(frame.data for frame in frames) if decode: return data.decode() else: return data @overload def get_iter(self, decode: Literal[True]) -> Iterator[str]: ... @overload def get_iter(self, decode: Literal[False]) -> Iterator[bytes]: ... @overload def get_iter(self, decode: bool | None = None) -> Iterator[Data]: ... def get_iter(self, decode: bool | None = None) -> Iterator[Data]: """ Stream the next message. Iterating the return value of :meth:`get_iter` yields a :class:`str` or :class:`bytes` for each frame in the message. The iterator must be fully consumed before calling :meth:`get_iter` or :meth:`get` again. Else, :exc:`ConcurrencyError` is raised. This method only makes sense for fragmented messages. If messages aren't fragmented, use :meth:`get` instead. Args: decode: :obj:`False` disables UTF-8 decoding of text frames and returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of binary frames and returns :class:`str`. Raises: EOFError: If the stream of frames has ended. UnicodeDecodeError: If a text frame contains invalid UTF-8. ConcurrencyError: If two coroutines run :meth:`get` or :meth:`get_iter` concurrently. """ with self.mutex: if self.get_in_progress: raise ConcurrencyError("get() or get_iter() is already running") self.get_in_progress = True # Locking with get_in_progress prevents concurrent execution # until get_iter() fetches a complete message or times out. # If get_iter() raises an exception e.g. in decoder.decode(), # get_in_progress remains set and the connection becomes unusable. # First frame frame = self.get_next_frame() with self.mutex: self.maybe_resume() assert frame.opcode is OP_TEXT or frame.opcode is OP_BINARY if decode is None: decode = frame.opcode is OP_TEXT if decode: decoder = UTF8Decoder() yield decoder.decode(frame.data, frame.fin) else: yield frame.data # Following frames, for fragmented messages while not frame.fin: frame = self.get_next_frame() with self.mutex: self.maybe_resume() assert frame.opcode is OP_CONT if decode: yield decoder.decode(frame.data, frame.fin) else: yield frame.data self.get_in_progress = False def put(self, frame: Frame) -> None: """ Add ``frame`` to the next message. Raises: EOFError: If the stream of frames has ended. """ with self.mutex: if self.closed: raise EOFError("stream of frames ended") self.frames.put(frame) self.maybe_pause() # put() and get/get_iter() call maybe_pause() and maybe_resume() while # holding self.mutex. This guarantees that the calls interleave properly. # Specifically, it prevents a race condition where maybe_resume() would # run before maybe_pause(), leaving the connection incorrectly paused. # A race condition is possible when get/get_iter() call self.frames.get() # without holding self.mutex. However, it's harmless — and even beneficial! # It can only result in popping an item from the queue before maybe_resume() # runs and skipping a pause() - resume() cycle that would otherwise occur. def maybe_pause(self) -> None: """Pause the writer if queue is above the high water mark.""" # Skip if flow control is disabled if self.high is None: return assert self.mutex.locked() # Check for "> high" to support high = 0 if self.frames.qsize() > self.high and not self.paused: self.paused = True self.pause() def maybe_resume(self) -> None: """Resume the writer if queue is below the low water mark.""" # Skip if flow control is disabled if self.low is None: return assert self.mutex.locked() # Check for "<= low" to support low = 0 if self.frames.qsize() <= self.low and self.paused: self.paused = False self.resume() def close(self) -> None: """ End the stream of frames. Calling :meth:`close` concurrently with :meth:`get`, :meth:`get_iter`, or :meth:`put` is safe. They will raise :exc:`EOFError`. """ with self.mutex: if self.closed: return self.closed = True if self.get_in_progress: # Unblock get() or get_iter(). self.frames.put(None) if self.paused: # Unblock recv_events(). self.paused = False self.resume()