second commit

env/lib/python3.11/site-packages/websockets/sync/messages.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import codecs
+import queue
+import threading
+from typing import Any, Callable, Iterable, Iterator
+
+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:
+                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)
+
+    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())
+            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())
+                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
+
+    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.
+
+        Callling :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)