Coverage for src/evutils/io/_source.py: 96%

1"""Byte-level input sources for event decoders. 

2 

3A :class:`ByteSource` abstracts *where* raw bytes come from (a file, an 

4in-memory buffer, a memory-mapped file, or -- in the future -- a live USB 

5device) from *how* they are parsed into events (the decoder). Decoders never 

6open files or ``seek`` raw handles; they only talk to a ByteSource. 

7 

8Two capabilities matter: 

9 

10* Every source supports sequential :meth:`ByteSource.read` -- the lowest common 

11 denominator, and the only thing a streaming device (USB, pipe) can do. 

12* Some sources are *mappable*: :meth:`ByteSource.buffer` hands out a zero-copy 

13 ``memoryview`` of the whole content (mmap, in-memory bytes). Decoders built on 

14 the native parser prefer this -- the C parser walks the bytes in place with no 

15 copy and no chunk-boundary carry logic. 

16 

17Lifetime note: a zero-copy ``buffer()`` (and any ``np.frombuffer`` view of it) 

18aliases the underlying storage. Drop those views *before* calling 

19:meth:`ByteSource.close`, or closing an mmap will raise ``BufferError``. 

20""" 

21from __future__ import annotations 

22 

23import io 

24import mmap 

25from abc import ABC, abstractmethod 

26from pathlib import Path 

27from typing import Any 

28 

29 

30class ByteSource(ABC): 

31 """Abstract raw-byte input. Knows nothing about events.""" 

32 

33 #: Filename if the source has one -- used for extension-based dispatch. 

34 name: str | None = None 

35 

36 @abstractmethod 

37 def read(self, size: int = -1) -> bytes: 

38 """Read up to ``size`` bytes (all remaining if ``size < 0``). 

39 

40 Returns ``b""`` at EOF. 

41 """ 

42 

43 @abstractmethod 

44 def peek(self, size: int) -> bytes: 

45 """Return up to ``size`` upcoming bytes *without* consuming them. 

46 

47 Used for header/magic sniffing. Should work even on non-seekable 

48 sources; raises :class:`io.UnsupportedOperation` when it cannot. 

49 """ 

50 

51 def readline(self) -> bytes: 

52 """Read one line, up to and including the newline (generic, byte-wise). 

53 

54 Used by text decoders (csv/txt) for header handling. Subclasses backed 

55 by a real stream override this with the stream's own ``readline``. 

56 """ 

57 out = bytearray() 

58 while True: 

59 b = self.read(1) 

60 if not b: 

61 break 

62 out += b 

63 if b == b"\n": 

64 break 

65 return bytes(out) 

66 

67 # -- optional zero-copy capability ------------------------------------- # 

68 def mappable(self) -> bool: 

69 """Check if the source is mappable. 

70 

71 Returns 

72 ------- 

73 bool 

74 True if mappable, False otherwise. 

75 

76 """ 

77 return False 

78 

79 def buffer(self) -> memoryview: 

80 """Zero-copy view of the *entire* content. Only if :meth:`mappable`.""" 

81 raise io.UnsupportedOperation("source is not mappable") 

82 

83 # -- optional random access -------------------------------------------- # 

84 def seekable(self) -> bool: 

85 """Check if the source is seekable. 

86 

87 Returns 

88 ------- 

89 bool 

90 True if seekable, False otherwise. 

91 

92 """ 

93 return False 

94 

95 def seek(self, pos: int, whence: int = io.SEEK_SET) -> int: 

96 """Seek to a specific position. 

97 

98 Parameters 

99 ---------- 

100 pos : int 

101 Position to seek to. 

102 whence : int, optional 

103 Reference point for seeking, by default io.SEEK_SET. 

104 

105 Returns 

106 ------- 

107 int 

108 New position. 

109 

110 """ 

111 raise io.UnsupportedOperation("source is not seekable") 

112 

113 def tell(self) -> int: 

114 """Get the current position. 

115 

116 Returns 

117 ------- 

118 int 

119 Current position. 

120 

121 """ 

122 raise io.UnsupportedOperation("source is not tellable") 

123 

124 def reset(self) -> None: 

125 """Return to the beginning of the stream.""" 

126 self.seek(0) 

127 

128 def close(self) -> None: 

129 """Close the source. 

130 

131 Returns 

132 ------- 

133 None 

134 

135 """ 

136 pass 

137 

138 def __enter__(self) -> "ByteSource": 

139 return self 

140 

141 def __exit__(self, *exc: Any) -> None: 

142 self.close() 

143 

144 

145def _clamp_seek(pos: int, whence: int, cur: int, length: int) -> int: 

146 if whence == io.SEEK_SET: 

147 new = pos 

148 elif whence == io.SEEK_CUR: 

149 new = cur + pos 

150 elif whence == io.SEEK_END: 

151 new = length + pos 

152 else: 

153 raise ValueError(f"invalid whence {whence}") 

154 return max(0, min(new, length)) 

155 

156 

157class StreamSource(ByteSource): 

158 """Wrap any binary stream exposing ``read`` (BufferedReader, pipe, ...). 

159 

160 This is the streaming fallback: it never claims to be mappable, so decoders 

161 read from it sequentially. 

162 """ 

163 

164 def __init__(self, stream: Any, name: str | None = None, owns: bool = False) -> None: 

165 if not hasattr(stream, "read"): 

166 raise TypeError("stream must have a read() method") 

167 self._s = stream 

168 n = name if name is not None else getattr(stream, "name", None) 

169 self.name = n if isinstance(n, str) else None 

170 self._owns = owns 

171 

172 def read(self, size: int = -1) -> bytes: 

173 return bytes(self._s.read(size)) 

174 

175 def readline(self) -> bytes: 

176 if hasattr(self._s, "readline"): 

177 return bytes(self._s.readline()) 

178 return super().readline() 

179 

180 def peek(self, size: int) -> bytes: 

181 s = self._s 

182 if hasattr(s, "peek"): 

183 return bytes(s.peek(size))[:size] 

184 if s.seekable(): 

185 pos = s.tell() 

186 try: 

187 return bytes(s.read(size)) 

188 finally: 

189 s.seek(pos) 

190 raise io.UnsupportedOperation( 

191 "source is not peekable; cannot sniff format -- pass an explicit decoder" 

192 ) 

193 

194 def seekable(self) -> bool: 

195 return bool(self._s.seekable()) 

196 

197 def seek(self, pos: int, whence: int = io.SEEK_SET) -> int: 

198 return int(self._s.seek(pos, whence)) 

199 

200 def tell(self) -> int: 

201 return int(self._s.tell()) 

202 

203 def close(self) -> None: 

204 if self._owns: 

205 self._s.close() 

206 

207 

208class BufferSource(ByteSource): 

209 """Zero-copy source over an in-memory buffer (bytes / bytearray / memoryview 

210 / ``BytesIO.getbuffer()``). 

211 """ 

212 

213 def __init__(self, data: Any, name: str | None = None) -> None: 

214 self._mv = memoryview(data).cast("B") # 1-D uint8 view, no copy 

215 self.name = name 

216 self._pos = 0 

217 

218 def read(self, size: int = -1) -> bytes: 

219 if size < 0: 

220 size = len(self._mv) - self._pos 

221 chunk = self._mv[self._pos:self._pos + size] 

222 self._pos += len(chunk) 

223 return bytes(chunk) 

224 

225 def peek(self, size: int) -> bytes: 

226 return bytes(self._mv[self._pos:self._pos + size]) 

227 

228 def mappable(self) -> bool: 

229 return True 

230 

231 def buffer(self) -> memoryview: 

232 return self._mv 

233 

234 def seekable(self) -> bool: 

235 return True 

236 

237 def seek(self, pos: int, whence: int = io.SEEK_SET) -> int: 

238 self._pos = _clamp_seek(pos, whence, self._pos, len(self._mv)) 

239 return self._pos 

240 

241 def tell(self) -> int: 

242 return self._pos 

243 

244 

245class MmapSource(ByteSource): 

246 """Memory-mapped, read-only file source. Zero-copy: the whole file is 

247 addressable, so the native parser can walk it in place. 

248 

249 Drop any ``buffer()``/``np.frombuffer`` views before :meth:`close`. 

250 """ 

251 

252 def __init__(self, path: str | Path) -> None: 

253 path = Path(path) 

254 self._f = open(path, "rb") 

255 try: 

256 self._mm = mmap.mmap(self._f.fileno(), 0, access=mmap.ACCESS_READ) 

257 except (ValueError, OSError): 

258 self._f.close() 

259 raise # empty file / unmappable -- caller falls back to StreamSource 

260 self.name = path.name 

261 self._pos = 0 

262 

263 def read(self, size: int = -1) -> bytes: 

264 if size < 0: 

265 size = len(self._mm) - self._pos 

266 chunk = self._mm[self._pos:self._pos + size] 

267 self._pos += len(chunk) 

268 return chunk 

269 

270 def peek(self, size: int) -> bytes: 

271 return bytes(self._mm[self._pos:self._pos + size]) 

272 

273 def mappable(self) -> bool: 

274 return True 

275 

276 def buffer(self) -> memoryview: 

277 return memoryview(self._mm) 

278 

279 def seekable(self) -> bool: 

280 return True 

281 

282 def seek(self, pos: int, whence: int = io.SEEK_SET) -> int: 

283 self._pos = _clamp_seek(pos, whence, self._pos, len(self._mm)) 

284 return self._pos 

285 

286 def tell(self) -> int: 

287 return self._pos 

288 

289 def close(self) -> None: 

290 # Raises BufferError if a caller still holds a buffer()/frombuffer view. 

291 self._mm.close() 

292 self._f.close() 

293 

294 

295def make_source(inp: Any, *, mmap_files: bool = True) -> ByteSource: 

296 """Normalise ``inp`` into a :class:`ByteSource`. 

297 

298 Accepts a path (str/Path), an in-memory buffer (bytes/bytearray/memoryview), 

299 a ``BytesIO``, any object with a binary ``read`` (e.g. ``BufferedReader``), 

300 or an already-constructed :class:`ByteSource` (returned as-is). 

301 

302 Regular files are memory-mapped by default (zero-copy); set 

303 ``mmap_files=False`` or fall back automatically for empty/unmappable files. 

304 """ 

305 if isinstance(inp, ByteSource): 

306 return inp 

307 if isinstance(inp, (str, Path)): 

308 p = Path(inp) 

309 if not p.is_file(): 

310 raise FileNotFoundError(f"File {p} does not exist") 

311 if mmap_files and p.stat().st_size > 0: 

312 try: 

313 return MmapSource(p) 

314 except (ValueError, OSError): 

315 pass 

316 return StreamSource(open(p, "rb"), name=p.name, owns=True) 

317 if isinstance(inp, (bytes, bytearray, memoryview)): 

318 return BufferSource(inp) 

319 if isinstance(inp, io.BytesIO): 

320 return BufferSource(inp.getbuffer()) # zero-copy view into the BytesIO 

321 if hasattr(inp, "read"): 

322 return StreamSource(inp) 

323 raise TypeError(f"Cannot create a ByteSource from {type(inp).__name__}")