Add trio message assembler.

Also uniformize code & tests with other implementations.

Author: aaugustin

src/websockets/asyncio/messages.py

@@ -35,7 +35,7 @@ def __len__(self) -> int:
  return len(self.queue)
 
  def put(self, item: T) -> None:
- """Put an item into the queue without waiting."""
+ """Put an item into the queue."""
  self.queue.append(item)
  if self.get_waiter is not None and not self.get_waiter.done():
  self.get_waiter.set_result(None)
@@ -81,8 +81,7 @@ class Assembler:
 
  """
 
- # coverage reports incorrectly: "line NN didn't jump to the function exit"
- def __init__( # pragma: no cover
+ def __init__(
  self,
  high: int | None = None,
  low: int | None = None,
@@ -155,15 +154,15 @@ async def get(self, decode: bool | None = None) -> Data:
  # until get() fetches a complete message or is canceled.
 
  try:
- # First frame
+ # Fetch the 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
+ # Fetch subsequent frames for fragmented messages.
  while not frame.fin:
  try:
  frame = await self.frames.get(not self.closed)
@@ -230,7 +229,7 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
  # If get_iter() raises an exception e.g. in decoder.decode(),
  # get_in_progress remains set and the connection becomes unusable.
 
- # First frame
+ # Yield the first frame.
  try:
  frame = await self.frames.get(not self.closed)
  except asyncio.CancelledError:
@@ -247,7 +246,7 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
  # Convert to bytes when frame.data is a bytearray.
  yield bytes(frame.data)
 
- # Following frames, for fragmented messages
+ # Yield subsequent 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()
@@ -280,22 +279,22 @@ def put(self, frame: Frame) -> None:
 
  def maybe_pause(self) -> None:
  """Pause the writer if queue is above the high water mark."""
- # Skip if flow control is disabled
+ # Skip if flow control is disabled.
  if self.high is None:
  return
 
- # Check for "> high" to support high = 0
+ # 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
+ # Skip if flow control is disabled.
  if self.low is None:
  return
 
- # Check for "<= low" to support low = 0
+ # Check for "<= low" to support low = 0.
  if len(self.frames) <= self.low and self.paused:
  self.paused = False
  self.resume()

src/websockets/sync/messages.py

@@ -165,7 +165,7 @@ def get(self, timeout: float | None = None, decode: bool | None = None) -> Data:
  try:
  deadline = Deadline(timeout)
 
- # First frame
+ # Fetch the first frame.
  frame = self.get_next_frame(deadline.timeout(raise_if_elapsed=False))
  with self.mutex:
  self.maybe_resume()
@@ -174,7 +174,7 @@ def get(self, timeout: float | None = None, decode: bool | None = None) -> Data:
  decode = frame.opcode is OP_TEXT
  frames = [frame]
 
- # Following frames, for fragmented messages
+ # Fetch subsequent frames for fragmented messages.
  while not frame.fin:
  try:
  frame = self.get_next_frame(
@@ -245,7 +245,7 @@ def get_iter(self, decode: bool | None = None) -> Iterator[Data]:
  # If get_iter() raises an exception e.g. in decoder.decode(),
  # get_in_progress remains set and the connection becomes unusable.
 
- # First frame
+ # Yield the first frame.
  frame = self.get_next_frame()
  with self.mutex:
  self.maybe_resume()
@@ -259,7 +259,7 @@ def get_iter(self, decode: bool | None = None) -> Iterator[Data]:
  # Convert to bytes when frame.data is a bytearray.
  yield bytes(frame.data)
 
- # Following frames, for fragmented messages
+ # Yield subsequent frames for fragmented messages.
  while not frame.fin:
  frame = self.get_next_frame()
  with self.mutex:
@@ -300,26 +300,26 @@ def put(self, frame: Frame) -> None:
 
  def maybe_pause(self) -> None:
  """Pause the writer if queue is above the high water mark."""
- # Skip if flow control is disabled
+ # Skip if flow control is disabled.
  if self.high is None:
  return
 
  assert self.mutex.locked()
 
- # Check for "> high" to support high = 0
+ # 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
+ # Skip if flow control is disabled.
  if self.low is None:
  return
 
  assert self.mutex.locked()
 
- # Check for "<= low" to support low = 0
+ # Check for "<= low" to support low = 0.
  if self.frames.qsize() <= self.low and self.paused:
  self.paused = False
  self.resume()