from __future__ import annotations from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import Any from .. import ClosedResourceError, DelimiterNotFound, EndOfStream, IncompleteRead from ..abc import AnyByteReceiveStream, ByteReceiveStream @dataclass(eq=False) class BufferedByteReceiveStream(ByteReceiveStream): """ Wraps any bytes-based receive stream and uses a buffer to provide sophisticated receiving capabilities in the form of a byte stream. """ receive_stream: AnyByteReceiveStream _buffer: bytearray = field(init=False, default_factory=bytearray) _closed: bool = field(init=False, default=False) async def aclose(self) -> None: await self.receive_stream.aclose() self._closed = True @property def buffer(self) -> bytes: """The bytes currently in the buffer.""" return bytes(self._buffer) @property def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]: return self.receive_stream.extra_attributes async def receive(self, max_bytes: int = 65536) -> bytes: if self._closed: raise ClosedResourceError if self._buffer: chunk = bytes(self._buffer[:max_bytes]) del self._buffer[:max_bytes] return chunk elif isinstance(self.receive_stream, ByteReceiveStream): return await self.receive_stream.receive(max_bytes) else: # With a bytes-oriented object stream, we need to handle any surplus bytes # we get from the receive() call chunk = await self.receive_stream.receive() if len(chunk) > max_bytes: # Save the surplus bytes in the buffer self._buffer.extend(chunk[max_bytes:]) return chunk[:max_bytes] else: return chunk async def receive_exactly(self, nbytes: int) -> bytes: """ Read exactly the given amount of bytes from the stream. :param nbytes: the number of bytes to read :return: the bytes read :raises ~anyio.IncompleteRead: if the stream was closed before the requested amount of bytes could be read from the stream """ while True: remaining = nbytes - len(self._buffer) if remaining <= 0: retval = self._buffer[:nbytes] del self._buffer[:nbytes] return bytes(retval) try: if isinstance(self.receive_stream, ByteReceiveStream): chunk = await self.receive_stream.receive(remaining) else: chunk = await self.receive_stream.receive() except EndOfStream as exc: raise IncompleteRead from exc self._buffer.extend(chunk) async def receive_until(self, delimiter: bytes, max_bytes: int) -> bytes: """ Read from the stream until the delimiter is found or max_bytes have been read. :param delimiter: the marker to look for in the stream :param max_bytes: maximum number of bytes that will be read before raising :exc:`~anyio.DelimiterNotFound` :return: the bytes read (not including the delimiter) :raises ~anyio.IncompleteRead: if the stream was closed before the delimiter was found :raises ~anyio.DelimiterNotFound: if the delimiter is not found within the bytes read up to the maximum allowed """ delimiter_size = len(delimiter) offset = 0 while True: # Check if the delimiter can be found in the current buffer index = self._buffer.find(delimiter, offset) if index >= 0: found = self._buffer[:index] del self._buffer[: index + len(delimiter) :] return bytes(found) # Check if the buffer is already at or over the limit if len(self._buffer) >= max_bytes: raise DelimiterNotFound(max_bytes) # Read more data into the buffer from the socket try: data = await self.receive_stream.receive() except EndOfStream as exc: raise IncompleteRead from exc # Move the offset forward and add the new data to the buffer offset = max(len(self._buffer) - delimiter_size + 1, 0) self._buffer.extend(data)