1  from __future__ import annotations
   2  
   3  from collections.abc import Callable, Mapping
   4  from dataclasses import dataclass, field
   5  from typing import Any
   6  
   7  from .. import ClosedResourceError, DelimiterNotFound, EndOfStream, IncompleteRead
   8  from ..abc import AnyByteReceiveStream, ByteReceiveStream
   9  
  10  
  11  @dataclass(eq=False)
  12  class BufferedByteReceiveStream(ByteReceiveStream):
  13      """
  14      Wraps any bytes-based receive stream and uses a buffer to provide sophisticated
  15      receiving capabilities in the form of a byte stream.
  16      """
  17  
  18      receive_stream: AnyByteReceiveStream
  19      _buffer: bytearray = field(init=False, default_factory=bytearray)
  20      _closed: bool = field(init=False, default=False)
  21  
  22      async def aclose(self) -> None:
  23          await self.receive_stream.aclose()
  24          self._closed = True
  25  
  26      @property
  27      def buffer(self) -> bytes:
  28          """The bytes currently in the buffer."""
  29          return bytes(self._buffer)
  30  
  31      @property
  32      def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  33          return self.receive_stream.extra_attributes
  34  
  35      async def receive(self, max_bytes: int = 65536) -> bytes:
  36          if self._closed:
  37              raise ClosedResourceError
  38  
  39          if self._buffer:
  40              chunk = bytes(self._buffer[:max_bytes])
  41              del self._buffer[:max_bytes]
  42              return chunk
  43          elif isinstance(self.receive_stream, ByteReceiveStream):
  44              return await self.receive_stream.receive(max_bytes)
  45          else:
  46              # With a bytes-oriented object stream, we need to handle any surplus bytes
  47              # we get from the receive() call
  48              chunk = await self.receive_stream.receive()
  49              if len(chunk) > max_bytes:
  50                  # Save the surplus bytes in the buffer
  51                  self._buffer.extend(chunk[max_bytes:])
  52                  return chunk[:max_bytes]
  53              else:
  54                  return chunk
  55  
  56      async def receive_exactly(self, nbytes: int) -> bytes:
  57          """
  58          Read exactly the given amount of bytes from the stream.
  59  
  60          :param nbytes: the number of bytes to read
  61          :return: the bytes read
  62          :raises ~anyio.IncompleteRead: if the stream was closed before the requested
  63              amount of bytes could be read from the stream
  64  
  65          """
  66          while True:
  67              remaining = nbytes - len(self._buffer)
  68              if remaining <= 0:
  69                  retval = self._buffer[:nbytes]
  70                  del self._buffer[:nbytes]
  71                  return bytes(retval)
  72  
  73              try:
  74                  if isinstance(self.receive_stream, ByteReceiveStream):
  75                      chunk = await self.receive_stream.receive(remaining)
  76                  else:
  77                      chunk = await self.receive_stream.receive()
  78              except EndOfStream as exc:
  79                  raise IncompleteRead from exc
  80  
  81              self._buffer.extend(chunk)
  82  
  83      async def receive_until(self, delimiter: bytes, max_bytes: int) -> bytes:
  84          """
  85          Read from the stream until the delimiter is found or max_bytes have been read.
  86  
  87          :param delimiter: the marker to look for in the stream
  88          :param max_bytes: maximum number of bytes that will be read before raising
  89              :exc:`~anyio.DelimiterNotFound`
  90          :return: the bytes read (not including the delimiter)
  91          :raises ~anyio.IncompleteRead: if the stream was closed before the delimiter
  92              was found
  93          :raises ~anyio.DelimiterNotFound: if the delimiter is not found within the
  94              bytes read up to the maximum allowed
  95  
  96          """
  97          delimiter_size = len(delimiter)
  98          offset = 0
  99          while True:
 100              # Check if the delimiter can be found in the current buffer
 101              index = self._buffer.find(delimiter, offset)
 102              if index >= 0:
 103                  found = self._buffer[:index]
 104                  del self._buffer[: index + len(delimiter) :]
 105                  return bytes(found)
 106  
 107              # Check if the buffer is already at or over the limit
 108              if len(self._buffer) >= max_bytes:
 109                  raise DelimiterNotFound(max_bytes)
 110  
 111              # Read more data into the buffer from the socket
 112              try:
 113                  data = await self.receive_stream.receive()
 114              except EndOfStream as exc:
 115                  raise IncompleteRead from exc
 116  
 117              # Move the offset forward and add the new data to the buffer
 118              offset = max(len(self._buffer) - delimiter_size + 1, 0)
 119              self._buffer.extend(data)
 120