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 

27 

28from ._compression import is_compressed_path, open_compressed 

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: object) -> None: 

142 self.close() 

143 

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

145 if whence == io.SEEK_SET: 

146 new = pos 

147 elif whence == io.SEEK_CUR: 

148 new = cur + pos 

149 elif whence == io.SEEK_END: 

150 new = length + pos 

151 else: 

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

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

154 

155class StreamSource(ByteSource): 

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

157 

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

159 read from it sequentially. 

160 """ 

161 

162 def __init__(self, stream: "io.BufferedIOBase", name: str | None = None, owns: bool = False) -> None: 

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

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

165 self._s = stream 

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

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

168 self._owns = owns 

169 

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

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

172 

173 def readline(self) -> bytes: 

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

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

176 return super().readline() 

177 

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

179 s = self._s 

180 if hasattr(s, "peek"): 

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

182 if s.seekable(): 

183 pos = s.tell() 

184 try: 

185 return bytes(s.read(size)) 

186 finally: 

187 s.seek(pos) 

188 raise io.UnsupportedOperation( 

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

190 ) 

191 

192 def seekable(self) -> bool: 

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

194 

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

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

197 

198 def tell(self) -> int: 

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

200 

201 def close(self) -> None: 

202 if self._owns: 

203 self._s.close() 

204 

205class BufferSource(ByteSource): 

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

207 / ``BytesIO.getbuffer()``). 

208 """ 

209 

210 def __init__(self, data: bytes, name: str | None = None) -> None: 

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

212 self.name = name 

213 self._pos = 0 

214 

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

216 if size < 0: 

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

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

219 self._pos += len(chunk) 

220 return bytes(chunk) 

221 

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

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

224 

225 def mappable(self) -> bool: 

226 return True 

227 

228 def buffer(self) -> memoryview: 

229 return self._mv 

230 

231 def seekable(self) -> bool: 

232 return True 

233 

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

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

236 return self._pos 

237 

238 def tell(self) -> int: 

239 return self._pos 

240 

241class MmapSource(ByteSource): 

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

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

244 

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

246 """ 

247 

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

249 path = Path(path) 

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

251 try: 

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

253 except (ValueError, OSError): 

254 self._f.close() 

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

256 self.name = path.name 

257 self._pos = 0 

258 

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

260 if size < 0: 

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

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

263 self._pos += len(chunk) 

264 return chunk 

265 

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

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

268 

269 def mappable(self) -> bool: 

270 return True 

271 

272 def buffer(self) -> memoryview: 

273 return memoryview(self._mm) 

274 

275 def seekable(self) -> bool: 

276 return True 

277 

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

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

280 return self._pos 

281 

282 def tell(self) -> int: 

283 return self._pos 

284 

285 def close(self) -> None: 

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

287 self._mm.close() 

288 self._f.close() 

289 

290def make_source(inp: "Path | str | bytes | io.BufferedIOBase", *, mmap_files: bool = True) -> ByteSource: 

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

292 

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

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

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

296 

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

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

299 """ 

300 if isinstance(inp, ByteSource): 

301 return inp 

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

303 p = Path(inp) 

304 if not p.is_file(): 

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

306 if is_compressed_path(p): 

307 # A compressed file cannot be mmap'd; wrap the decompressing stream. 

308 # Keep the full name (incl. the compression suffix) so format 

309 # detection can strip it back to the inner extension. 

310 return StreamSource(open_compressed(p, "rb"), name=p.name, owns=True) 

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__}")