Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/streams/buffered.py: 37%

1from __future__ import annotations 

2 

3__all__ = ( 

4 "BufferedByteReceiveStream", 

5 "BufferedByteStream", 

6 "BufferedConnectable", 

7) 

8 

9import sys 

10from collections.abc import Callable, Iterable, Mapping 

11from dataclasses import dataclass, field 

12from typing import Any, SupportsIndex 

13 

14from .. import ClosedResourceError, DelimiterNotFound, EndOfStream, IncompleteRead 

15from ..abc import ( 

16 AnyByteReceiveStream, 

17 AnyByteStream, 

18 AnyByteStreamConnectable, 

19 ByteReceiveStream, 

20 ByteStream, 

21 ByteStreamConnectable, 

22) 

23 

24if sys.version_info >= (3, 12): 

25 from typing import override 

26else: 

27 from typing_extensions import override 

28 

29 

30@dataclass(eq=False) 

31class BufferedByteReceiveStream(ByteReceiveStream): 

32 """ 

33 Wraps any bytes-based receive stream and uses a buffer to provide sophisticated 

34 receiving capabilities in the form of a byte stream. 

35 """ 

36 

37 receive_stream: AnyByteReceiveStream 

38 _buffer: bytearray = field(init=False, default_factory=bytearray) 

39 _closed: bool = field(init=False, default=False) 

40 

41 async def aclose(self) -> None: 

42 await self.receive_stream.aclose() 

43 self._closed = True 

44 

45 @property 

46 def buffer(self) -> bytes: 

47 """The bytes currently in the buffer.""" 

48 return bytes(self._buffer) 

49 

50 @property 

51 def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]: 

52 return self.receive_stream.extra_attributes 

53 

54 def feed_data(self, data: Iterable[SupportsIndex], /) -> None: 

55 """ 

56 Append data directly into the buffer. 

57 

58 Any data in the buffer will be consumed by receive operations before receiving 

59 anything from the wrapped stream. 

60 

61 :param data: the data to append to the buffer (can be bytes or anything else 

62 that supports ``__index__()``) 

63 

64 """ 

65 self._buffer.extend(data) 

66 

67 async def receive(self, max_bytes: int = 65536) -> bytes: 

68 if max_bytes < 1: 

69 raise ValueError("max_bytes must be a positive integer") 

70 

71 if self._closed: 

72 raise ClosedResourceError 

73 

74 if self._buffer: 

75 chunk = bytes(self._buffer[:max_bytes]) 

76 del self._buffer[:max_bytes] 

77 return chunk 

78 elif isinstance(self.receive_stream, ByteReceiveStream): 

79 return await self.receive_stream.receive(max_bytes) 

80 else: 

81 # With a bytes-oriented object stream, we need to handle any surplus bytes 

82 # we get from the receive() call 

83 chunk = await self.receive_stream.receive() 

84 if len(chunk) > max_bytes: 

85 # Save the surplus bytes in the buffer 

86 self._buffer.extend(chunk[max_bytes:]) 

87 return chunk[:max_bytes] 

88 else: 

89 return chunk 

90 

91 async def receive_exactly(self, nbytes: int) -> bytes: 

92 """ 

93 Read exactly the given amount of bytes from the stream. 

94 

95 :param nbytes: the number of bytes to read 

96 :return: the bytes read 

97 :raises ~anyio.IncompleteRead: if the stream was closed before the requested 

98 amount of bytes could be read from the stream 

99 

100 """ 

101 while True: 

102 remaining = nbytes - len(self._buffer) 

103 if remaining <= 0: 

104 retval = self._buffer[:nbytes] 

105 del self._buffer[:nbytes] 

106 return bytes(retval) 

107 

108 try: 

109 if isinstance(self.receive_stream, ByteReceiveStream): 

110 chunk = await self.receive_stream.receive(remaining) 

111 else: 

112 chunk = await self.receive_stream.receive() 

113 except EndOfStream as exc: 

114 raise IncompleteRead from exc 

115 

116 self._buffer.extend(chunk) 

117 

118 async def receive_until(self, delimiter: bytes, max_bytes: int) -> bytes: 

119 """ 

120 Read from the stream until the delimiter is found or max_bytes have been read. 

121 

122 :param delimiter: the marker to look for in the stream 

123 :param max_bytes: maximum number of bytes that will be read before raising 

124 :exc:`~anyio.DelimiterNotFound` 

125 :return: the bytes read (not including the delimiter) 

126 :raises ~anyio.IncompleteRead: if the stream was closed before the delimiter 

127 was found 

128 :raises ~anyio.DelimiterNotFound: if the delimiter is not found within the 

129 bytes read up to the maximum allowed 

130 

131 """ 

132 delimiter_size = len(delimiter) 

133 offset = 0 

134 while True: 

135 # Check if the delimiter can be found in the current buffer 

136 index = self._buffer.find(delimiter, offset) 

137 if index >= 0: 

138 found = self._buffer[:index] 

139 del self._buffer[: index + len(delimiter) :] 

140 return bytes(found) 

141 

142 # Check if the buffer is already at or over the limit 

143 if len(self._buffer) >= max_bytes: 

144 raise DelimiterNotFound(max_bytes) 

145 

146 # Read more data into the buffer from the socket 

147 try: 

148 data = await self.receive_stream.receive() 

149 except EndOfStream as exc: 

150 raise IncompleteRead from exc 

151 

152 # Move the offset forward and add the new data to the buffer 

153 offset = max(len(self._buffer) - delimiter_size + 1, 0) 

154 self._buffer.extend(data) 

155 

156 

157class BufferedByteStream(BufferedByteReceiveStream, ByteStream): 

158 """ 

159 A full-duplex variant of :class:`BufferedByteReceiveStream`. All writes are passed 

160 through to the wrapped stream as-is. 

161 """ 

162 

163 def __init__(self, stream: AnyByteStream): 

164 """ 

165 :param stream: the stream to be wrapped 

166 

167 """ 

168 super().__init__(stream) 

169 self._stream = stream 

170 

171 @override 

172 async def send_eof(self) -> None: 

173 await self._stream.send_eof() 

174 

175 @override 

176 async def send(self, item: bytes) -> None: 

177 await self._stream.send(item) 

178 

179 

180class BufferedConnectable(ByteStreamConnectable): 

181 """ 

182 Wraps a byte stream connectable to produce :class:`BufferedByteStream` connections. 

183 

184 Use this when you want the streams returned by :meth:`connect` to have the buffered 

185 receive API (e.g. :meth:`~BufferedByteReceiveStream.receive_exactly` and 

186 :meth:`~BufferedByteReceiveStream.receive_until`). 

187 

188 :param connectable: the byte stream connectable to wrap 

189 """ 

190 

191 def __init__(self, connectable: AnyByteStreamConnectable): 

192 """ 

193 :param connectable: the connectable to wrap 

194 

195 """ 

196 self.connectable = connectable 

197 

198 @override 

199 async def connect(self) -> BufferedByteStream: 

200 stream = await self.connectable.connect() 

201 return BufferedByteStream(stream)