297 lines
10 KiB
Python
297 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import codecs
|
|
import collections
|
|
from collections.abc import AsyncIterator, Iterable
|
|
from typing import Any, Callable, Generic, TypeVar
|
|
|
|
from ..exceptions import ConcurrencyError
|
|
from ..frames import OP_BINARY, OP_CONT, OP_TEXT, Frame
|
|
from ..typing import Data
|
|
|
|
|
|
__all__ = ["Assembler"]
|
|
|
|
UTF8Decoder = codecs.getincrementaldecoder("utf-8")
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class SimpleQueue(Generic[T]):
|
|
"""
|
|
Simplified version of :class:`asyncio.Queue`.
|
|
|
|
Provides only the subset of functionality needed by :class:`Assembler`.
|
|
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.loop = asyncio.get_running_loop()
|
|
self.get_waiter: asyncio.Future[None] | None = None
|
|
self.queue: collections.deque[T] = collections.deque()
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.queue)
|
|
|
|
def put(self, item: T) -> None:
|
|
"""Put an item into the queue without waiting."""
|
|
self.queue.append(item)
|
|
if self.get_waiter is not None and not self.get_waiter.done():
|
|
self.get_waiter.set_result(None)
|
|
|
|
async def get(self, block: bool = True) -> T:
|
|
"""Remove and return an item from the queue, waiting if necessary."""
|
|
if not self.queue:
|
|
if not block:
|
|
raise EOFError("stream of frames ended")
|
|
assert self.get_waiter is None, "cannot call get() concurrently"
|
|
self.get_waiter = self.loop.create_future()
|
|
try:
|
|
await self.get_waiter
|
|
finally:
|
|
self.get_waiter.cancel()
|
|
self.get_waiter = None
|
|
return self.queue.popleft()
|
|
|
|
def reset(self, items: Iterable[T]) -> None:
|
|
"""Put back items into an empty, idle queue."""
|
|
assert self.get_waiter is None, "cannot reset() while get() is running"
|
|
assert not self.queue, "cannot reset() while queue isn't empty"
|
|
self.queue.extend(items)
|
|
|
|
def abort(self) -> None:
|
|
"""Close the queue, raising EOFError in get() if necessary."""
|
|
if self.get_waiter is not None and not self.get_waiter.done():
|
|
self.get_waiter.set_exception(EOFError("stream of frames ended"))
|
|
|
|
|
|
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.
|
|
|
|
"""
|
|
|
|
# coverage reports incorrectly: "line NN didn't jump to the function exit"
|
|
def __init__( # pragma: no cover
|
|
self,
|
|
high: int | None = None,
|
|
low: int | None = None,
|
|
pause: Callable[[], Any] = lambda: None,
|
|
resume: Callable[[], Any] = lambda: None,
|
|
) -> None:
|
|
# Queue of incoming frames.
|
|
self.frames: SimpleQueue[Frame] = 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
|
|
|
|
async def get(self, 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:
|
|
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.
|
|
|
|
"""
|
|
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 is cancelled.
|
|
|
|
try:
|
|
# First frame
|
|
frame = await self.frames.get(not self.closed)
|
|
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 = await self.frames.get(not self.closed)
|
|
except asyncio.CancelledError:
|
|
# Put frames already received back into the queue
|
|
# so that future calls to get() can return them.
|
|
self.frames.reset(frames)
|
|
raise
|
|
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
|
|
|
|
async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
|
|
"""
|
|
Stream the next message.
|
|
|
|
Iterating the return value of :meth:`get_iter` asynchronously 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.
|
|
|
|
"""
|
|
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 is cancelled.
|
|
|
|
# If get_iter() raises an exception e.g. in decoder.decode(),
|
|
# get_in_progress remains set and the connection becomes unusable.
|
|
|
|
# First frame
|
|
try:
|
|
frame = await self.frames.get(not self.closed)
|
|
except asyncio.CancelledError:
|
|
self.get_in_progress = False
|
|
raise
|
|
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:
|
|
# We cannot handle asyncio.CancelledError because we don't buffer
|
|
# previous fragments — we're streaming them. Canceling get_iter()
|
|
# here will leave the assembler in a stuck state. Future calls to
|
|
# get() or get_iter() will raise ConcurrencyError.
|
|
frame = await self.frames.get(not self.closed)
|
|
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.
|
|
|
|
"""
|
|
if self.closed:
|
|
raise EOFError("stream of frames ended")
|
|
|
|
self.frames.put(frame)
|
|
self.maybe_pause()
|
|
|
|
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
|
|
|
|
# Check for "> high" to support high = 0
|
|
if len(self.frames) > 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
|
|
|
|
# Check for "<= low" to support low = 0
|
|
if len(self.frames) <= 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`.
|
|
|
|
"""
|
|
if self.closed:
|
|
return
|
|
|
|
self.closed = True
|
|
|
|
# Unblock get() or get_iter().
|
|
self.frames.abort()
|