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

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 while True: 

84 if chunk := await self.receive_stream.receive(): 

85 break 

86 

87 if len(chunk) > max_bytes: 

88 # Save the surplus bytes in the buffer 

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

90 return chunk[:max_bytes] 

91 else: 

92 return chunk 

93 

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

95 """ 

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

97 

98 :param nbytes: the number of bytes to read 

99 :return: the bytes read 

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

101 amount of bytes could be read from the stream 

102 :raises ~anyio.ClosedResourceError: if this stream has been explicitly closed 

103 

104 """ 

105 if self._closed: 

106 raise ClosedResourceError 

107 

108 while True: 

109 remaining = nbytes - len(self._buffer) 

110 if remaining <= 0: 

111 retval = self._buffer[:nbytes] 

112 del self._buffer[:nbytes] 

113 return bytes(retval) 

114 

115 try: 

116 if isinstance(self.receive_stream, ByteReceiveStream): 

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

118 else: 

119 chunk = await self.receive_stream.receive() 

120 except EndOfStream as exc: 

121 raise IncompleteRead from exc 

122 

123 self._buffer.extend(chunk) 

124 

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

126 """ 

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

128 

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

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

131 :exc:`~anyio.DelimiterNotFound` 

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

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

134 was found 

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

136 bytes read up to the maximum allowed 

137 :raises ~anyio.ClosedResourceError: if this stream has been explicitly closed 

138 

139 """ 

140 if self._closed: 

141 raise ClosedResourceError 

142 

143 delimiter_size = len(delimiter) 

144 offset = 0 

145 while True: 

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

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

148 if index >= 0: 

149 found = self._buffer[:index] 

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

151 return bytes(found) 

152 

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

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

155 raise DelimiterNotFound(max_bytes) 

156 

157 # Read more data into the buffer from the socket 

158 try: 

159 data = await self.receive_stream.receive() 

160 except EndOfStream as exc: 

161 raise IncompleteRead from exc 

162 

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

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

165 self._buffer.extend(data) 

166 

167 

168class BufferedByteStream(BufferedByteReceiveStream, ByteStream): 

169 """ 

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

171 through to the wrapped stream as-is. 

172 """ 

173 

174 def __init__(self, stream: AnyByteStream): 

175 """ 

176 :param stream: the stream to be wrapped 

177 

178 """ 

179 super().__init__(stream) 

180 self._stream = stream 

181 

182 @override 

183 async def send_eof(self) -> None: 

184 await self._stream.send_eof() 

185 

186 @override 

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

188 await self._stream.send(item) 

189 

190 

191class BufferedConnectable(ByteStreamConnectable): 

192 """ 

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

194 

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

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

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

198 

199 :param connectable: the byte stream connectable to wrap 

200 """ 

201 

202 def __init__(self, connectable: AnyByteStreamConnectable): 

203 """ 

204 :param connectable: the connectable to wrap 

205 

206 """ 

207 self.connectable = connectable 

208 

209 @override 

210 async def connect(self) -> BufferedByteStream: 

211 stream = await self.connectable.connect() 

212 return BufferedByteStream(stream)