Coverage for core / src / sensorkit / data / streams.py: 78%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5from collections.abc import Buffer 

6from typing import Protocol, runtime_checkable 

7 

8type Buf = bytes | bytearray | memoryview 

9 

10 

11@runtime_checkable 

12class StreamReader(Protocol): 

13 """A variant of the asyncio.StreamReader interface.""" 

14 

15 def exception(self) -> Exception: ... 

16 def at_eof(self) -> bool: ... 

17 async def readline(self) -> Buf: ... 

18 async def readuntil(self, separator=b"\n") -> Buf: ... 

19 async def read(self, n=-1) -> Buf: ... 

20 async def readexactly(self, n: int) -> Buf: ... 

21 

22 async def __aiter__(self): 

23 while not self.at_eof(): 

24 yield await self.read() 

25 

26 

27@runtime_checkable 

28class StreamWriter(Protocol): 

29 """A variant of the asyncio.StreamWriter interface.""" 

30 

31 def write(self, data: bytes): ... 

32 def close(self): ... 

33 def is_closing(self) -> bool: ... 

34 async def wait_closed(self): ... 

35 async def drain(self): ... 

36 

37 

38class BufferReader(StreamReader): 

39 """A StreamReader that reads from an immutable input buffer.""" 

40 

41 def __init__(self, data: Buffer): 

42 self.data: Buf = memoryview(data) 

43 self.ptr = 0 

44 self._exception = None 

45 

46 def exception(self) -> Exception: 

47 return self._exception 

48 

49 def at_eof(self) -> bool: 

50 return self.ptr >= len(self.data) 

51 

52 async def readline(self) -> Buf: 

53 return await self.readuntil() 

54 

55 async def readuntil(self, separator=b"\n") -> Buf: 

56 try: 

57 start = self.ptr 

58 sep_len = len(separator) 

59 

60 # Find the separator in the buffer 

61 for i in range(start, len(self.data) - sep_len + 1): 

62 if self.data[i:i+sep_len] == separator: 

63 self.ptr = i + sep_len 

64 return self.data[start:self.ptr] 

65 

66 # Separator not found, return the rest of the buffer 

67 self.ptr = len(self.data) 

68 return self.data[start:] 

69 except Exception as e: 

70 self._exception = e 

71 raise 

72 

73 async def read(self, n=-1) -> Buf: 

74 try: 

75 start = self.ptr 

76 end = min(start + n, len(self.data)) if n >= 0 else len(self.data) 

77 self.ptr = end 

78 return self.data[start:end] 

79 except Exception as e: 

80 self._exception = e 

81 raise 

82 

83 async def readexactly(self, n: int) -> Buf: 

84 buf = await self.read(n) 

85 

86 if len(buf) != n: 

87 raise asyncio.IncompleteReadError(bytes(buf), n) 

88 

89 return buf 

90 

91 

92class BufferWriter(StreamWriter): 

93 """A StreamWriter that writes to a buffer.""" 

94 

95 def __init__(self, loop: asyncio.AbstractEventLoop | None = None): 

96 loop = loop or asyncio.get_event_loop() 

97 self._future: asyncio.Future[Buffer] = loop.create_future() 

98 self._bytes = bytearray() 

99 

100 def write(self, data: bytes): 

101 self._bytes.extend(data) 

102 

103 def write_eof(self): 

104 # No eof. 

105 pass 

106 

107 def close(self): 

108 self._future.set_result(self._bytes) 

109 

110 def is_closing(self) -> bool: 

111 return self._future.done() 

112 

113 async def wait_closed(self): 

114 await self._future 

115 

116 async def drain(self): 

117 # Writes are immediate. 

118 pass 

119 

120 def get_future(self): 

121 """Return a Future that resolves to the accumulated buffer once the writer is closed.""" 

122 return self._future 

123 

124 

125class QueueReader(StreamReader): 

126 """A StreamReader that reads from a queue with minimal copying.""" 

127 

128 def __init__(self, queue: asyncio.Queue[bytes | None]): 

129 self._queue = queue 

130 self._buffer = bytearray() # Only used for partial chunks 

131 self._eof = False 

132 self._exception: Exception | None = None 

133 

134 def exception(self) -> Exception: 

135 return self._exception 

136 

137 def at_eof(self) -> bool: 

138 return self._eof and len(self._buffer) == 0 

139 

140 async def _get_next_chunk(self): 

141 """Get next chunk from queue or return None on EOF.""" 

142 if self._eof: 

143 return None 

144 

145 try: 

146 chunk = await self._queue.get() 

147 self._queue.task_done() 

148 

149 # Check for end of stream marker. 

150 if chunk is None: 

151 self._eof = True 

152 return None 

153 

154 return chunk 

155 except Exception as e: 

156 self._exception = e 

157 raise 

158 

159 async def read(self, n=-1) -> Buf: 

160 if n == 0: 

161 return b"" 

162 

163 # If we have data in the buffer, handle that first 

164 if self._buffer: 

165 if n == -1: 

166 n = len(self._buffer) 

167 

168 if len(self._buffer) <= n: 

169 remainder = bytearray() 

170 else: 

171 remainder = self._buffer[n:] 

172 del self._buffer[n:] 

173 

174 result = self._buffer 

175 self._buffer = remainder 

176 return result 

177 

178 # No buffer data, read directly from queue 

179 chunk = await self._get_next_chunk() 

180 

181 if chunk is None: 

182 return b"" 

183 

184 if n == -1: 

185 n = len(chunk) 

186 

187 if len(chunk) <= n: 

188 return chunk 

189 

190 # Store excess in buffer 

191 view = memoryview(chunk) 

192 self._buffer = bytearray(view[n:]) 

193 return view[:n] 

194 

195 async def readline(self) -> Buf: 

196 return await self.readuntil() 

197 

198 async def readuntil(self, separator=b"\n") -> Buf: 

199 idx = 0 

200 

201 while True: 

202 # Check whether the newest part of the current buffer has the needle. 

203 try: 

204 idx = self._buffer.index(separator, idx) + len(separator) 

205 remainder = self._buffer[idx:] 

206 del self._buffer[idx:] 

207 result = self._buffer 

208 self._buffer = remainder 

209 return result 

210 except ValueError: 

211 idx = len(self._buffer) - len(separator) + 1 

212 

213 # Read the next chunk from the queue. 

214 chunk = await self._get_next_chunk() 

215 

216 if chunk is None: 

217 result = self._buffer 

218 self._buffer = bytearray() 

219 return result 

220 

221 self._buffer.extend(chunk) 

222 

223 async def readexactly(self, n: int) -> Buf: 

224 while len(self._buffer) < n: 

225 chunk = await self._get_next_chunk() 

226 

227 if chunk is None: 

228 partial = bytes(self._buffer) 

229 self._buffer = bytearray() 

230 raise asyncio.IncompleteReadError(partial, n) 

231 

232 self._buffer.extend(chunk) 

233 

234 remainder = self._buffer[n:] 

235 del self._buffer[n:] 

236 result = self._buffer 

237 self._buffer = remainder 

238 return result 

239 

240 

241class QueueWriter(StreamWriter): 

242 """A StreamWriter that writes to a queue.""" 

243 

244 def __init__(self, queue: asyncio.Queue[bytes | None]): 

245 self._queue = queue 

246 self._closed = False 

247 

248 def write(self, data: bytes): 

249 if self._closed: 

250 raise RuntimeError("Writer is closed") 

251 self._queue.put_nowait(data) 

252 

253 def close(self): 

254 if not self._closed: 

255 # Signal end of stream. 

256 self._queue.put_nowait(None) 

257 self._closed = True 

258 

259 def is_closing(self) -> bool: 

260 return self._closed 

261 

262 async def wait_closed(self): 

263 if self._closed: 

264 await self._queue.join() 

265 

266 async def drain(self): 

267 # Writes are immediate. 

268 pass 

269 

270 

271def create_connected_streams() -> tuple[StreamReader, StreamWriter]: 

272 """Return a pair of streams that are connected by a queue.""" 

273 queue = asyncio.Queue() 

274 return QueueReader(queue), QueueWriter(queue)