Coverage for src/evutils/io/_npz.py: 88%
230 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""NPZ file decoder and encoder.
3Events are stored as four flat arrays under the keys ``t``, ``x``, ``y`` and
4``p`` (the SoA layout of :class:`~evutils.types.EventArray`), plus optional
5scalar ``width`` / ``height`` entries. A single structured array under the key
6``events`` (:data:`~evutils.types.Event_dtype`-like) is also accepted when
7reading. The layout is fully compatible with plain
8``np.savez(f, t=..., x=..., y=..., p=...)`` / ``np.load``.
10Both directions stream and never materialise the whole recording:
12* The decoder reads the ``.npy`` members through zip streams chunk by chunk
13 (works for stored and deflated members alike).
14* The encoder cannot write four zip members simultaneously (the zip format is
15 strictly sequential), so :meth:`~EventEncoder_Npz.write` spools each column
16 to an unlinked temporary file as raw bytes; :meth:`~EventEncoder_Npz.close`
17 then streams every spool into the archive as a proper ``.npy`` member.
18"""
19from __future__ import annotations
21import io
22import tempfile
23import zipfile
24from datetime import datetime
25from typing import IO
27import numpy as np
28from numpy.lib import format as npy_format
30from ..types import EventArray, TriggerArray
31from .common import EventDecoder, EventEncoder
32from ._source import ByteSource
34_EMPTY_EVENTS = EventArray.empty()
36#: Column name -> on-disk dtype (matches EventArray's column dtypes).
37_COLUMNS = (("t", np.dtype(np.int64)), ("x", np.dtype(np.uint16)),
38 ("y", np.dtype(np.uint16)), ("p", np.dtype(np.uint8)))
40def _read_npy_header(fp: IO[bytes]) -> tuple[tuple[int, ...], np.dtype]:
41 """Read the ``.npy`` magic + header from a stream, returning (shape, dtype).
43 Leaves ``fp`` positioned at the first data byte. Fortran-ordered arrays are
44 rejected (event columns are 1-D, written C-ordered).
45 """
46 version = npy_format.read_magic(fp)
47 read_header = {
48 (1, 0): npy_format.read_array_header_1_0,
49 (2, 0): npy_format.read_array_header_2_0,
50 }.get(version)
51 if read_header is None:
52 raise ValueError(f"unsupported .npy format version {version}")
53 shape, fortran, dtype = read_header(fp)
54 if fortran:
55 raise ValueError("Fortran-ordered .npy members are not supported")
56 return shape, dtype
58def _read_exact(fp: IO[bytes], nbytes: int) -> bytearray:
59 """Read exactly ``nbytes`` from a (possibly decompressing) stream.
61 Returns a writable buffer so the numpy views over it are mutable.
62 """
63 out = bytearray(nbytes)
64 view = memoryview(out)
65 got = 0
66 while got < nbytes:
67 n = fp.readinto(view[got:]) # type: ignore[attr-defined]
68 if not n:
69 raise EOFError(f"truncated .npy member: expected {nbytes} bytes, got {got}")
70 got += n
71 return out
73class EventDecoder_Npz(EventDecoder):
74 """Decode events from an ``.npz`` archive, streaming chunk by chunk.
76 The archive members are read through zip streams: only ``chunk_size``
77 events are held in memory at a time, so arbitrarily large recordings can
78 be iterated. Accepts either the four column members ``t/x/y/p`` or a
79 single structured ``events`` member.
81 Parameters
82 ----------
83 source
84 Byte source to read from (must be seekable, as required by the zip
85 format).
86 chunk_size
87 Number of events returned per :meth:`read_chunk` call.
89 """
91 #: read_chunk returns fresh, independent arrays bounded by n_events_hint, so
92 #: EventReader can hand them out directly (skipping the staging accumulator).
93 _independent_windows = True
95 #: NPZ columns are index-addressable, so seeking is a stream reposition (by
96 #: event index) or a searchsorted over the timestamp column (by time).
97 SUPPORTS_SEEK = True
99 def __init__(self, source: ByteSource, chunk_size: int = 1_000_000):
100 super().__init__(source, chunk_size)
101 self._zf: zipfile.ZipFile | None = None
102 self._streams: dict[str, IO[bytes]] = {}
103 self._aos_dtype: np.dtype | None = None # set when reading an 'events' member
104 self._n = 0
105 self._pos = 0
107 def _open_streams(self) -> None:
108 """(Re)open the member streams and position them past the npy headers."""
109 assert self._zf is not None
110 for fp in self._streams.values():
111 fp.close()
112 self._streams = {}
114 names = set(self._zf.namelist())
115 if {"t.npy", "x.npy", "y.npy", "p.npy"} <= names:
116 n = None
117 for name, _ in _COLUMNS:
118 fp = self._zf.open(f"{name}.npy")
119 shape, dtype = _read_npy_header(fp)
120 if len(shape) != 1:
121 raise ValueError(f"member {name}.npy is not 1-D: shape {shape}")
122 if n is None:
123 n = shape[0]
124 elif shape[0] != n:
125 raise ValueError("event columns have mismatched lengths")
126 self._streams[name] = fp
127 self._n = n or 0
128 self._aos_dtype = None
129 elif "events.npy" in names:
130 fp = self._zf.open("events.npy")
131 shape, dtype = _read_npy_header(fp)
132 if dtype.names is None or not {"t", "x", "y", "p"} <= set(dtype.names):
133 raise ValueError("'events' member must be a structured array with t/x/y/p fields")
134 self._streams["events"] = fp
135 self._aos_dtype = dtype
136 self._n = shape[0]
137 else:
138 raise ValueError(
139 f"NPZ archive does not contain event data: expected members "
140 f"'t/x/y/p' or 'events', found {sorted(names)}"
141 )
143 def init(self) -> None:
144 """Open the archive and locate the event members."""
145 if self._is_initialized:
146 return
148 f: "io.BytesIO | io.BufferedIOBase" = self._source if self._source.seekable() else io.BytesIO(self._source.read(-1))
149 self._zf = zipfile.ZipFile(f)
151 names = set(self._zf.namelist())
152 for attr, member in (("_width", "width.npy"), ("_height", "height.npy")):
153 if member in names:
154 with self._zf.open(member) as fp:
155 setattr(self, attr, int(npy_format.read_array(fp).item()))
157 self._open_streams()
158 self._pos = 0
159 self._is_initialized = True
161 def _read_n(self, n: int) -> EventArray:
162 """Stream the next ``n`` events out of the member streams."""
163 if self._aos_dtype is not None:
164 fp = self._streams["events"]
165 buf = _read_exact(fp, n * self._aos_dtype.itemsize)
166 return EventArray.from_aos(np.frombuffer(buf, dtype=self._aos_dtype))
167 cols = {}
168 for name, dtype in _COLUMNS:
169 # The member's own dtype was validated against 1-D at open; event
170 # columns are cast to the canonical dtypes by EventArray.
171 buf = _read_exact(self._streams[name], n * dtype.itemsize)
172 cols[name] = np.frombuffer(buf, dtype=dtype)
173 return EventArray(cols["t"], cols["x"], cols["y"], cols["p"])
175 def read_chunk(self, delta_t_hint: int | None = None,
176 n_events_hint: int | None = None) -> EventArray:
177 if not self._is_initialized:
178 self.init()
180 if self._pos >= self._n:
181 self._eof = True
182 return _EMPTY_EVENTS
184 n = min(n_events_hint or self._chunk_size, self._n - self._pos)
185 chunk = self._read_n(n)
186 self._pos += n
187 if self._pos >= self._n:
188 self._eof = True
189 return chunk
191 def _itemsize(self, name: str) -> int:
192 if name == "events":
193 assert self._aos_dtype is not None
194 return self._aos_dtype.itemsize
195 return dict(_COLUMNS)[name].itemsize
197 def _load_all_t(self) -> np.ndarray:
198 """Read the full timestamp column once (for a time->index search)."""
199 assert self._zf is not None
200 if self._aos_dtype is not None:
201 with self._zf.open("events.npy") as fp:
202 return npy_format.read_array(fp)["t"]
203 with self._zf.open("t.npy") as fp:
204 return npy_format.read_array(fp)
206 def _ts_at(self, idx: int) -> int | None:
207 """Timestamp of event ``idx`` (``None`` if at/after the end)."""
208 assert self._zf is not None
209 if idx >= self._n:
210 return None
211 if self._aos_dtype is not None:
212 with self._zf.open("events.npy") as fp:
213 _read_npy_header(fp)
214 skip_bytes = idx * self._aos_dtype.itemsize
215 while skip_bytes > 0:
216 chunk = fp.read(min(skip_bytes, 1 << 20))
217 if not chunk:
218 break
219 skip_bytes -= len(chunk)
220 rec = np.frombuffer(_read_exact(fp, self._aos_dtype.itemsize),
221 dtype=self._aos_dtype)
222 return int(rec["t"][0])
223 with self._zf.open("t.npy") as fp:
224 _read_npy_header(fp)
225 skip_bytes = idx * 8
226 while skip_bytes > 0:
227 chunk = fp.read(min(skip_bytes, 1 << 20))
228 if not chunk:
229 break
230 skip_bytes -= len(chunk)
231 return int(np.frombuffer(_read_exact(fp, 8), dtype=np.int64)[0])
233 def _seek_to_index(self, idx: int) -> None:
234 """Reposition every member stream so the next read starts at ``idx``."""
235 self._open_streams() # streams sit at their first data byte (headers consumed)
236 for name, fp in self._streams.items():
237 skip_bytes = idx * self._itemsize(name)
238 while skip_bytes > 0:
239 chunk = fp.read(min(skip_bytes, 1 << 20))
240 if not chunk:
241 break
242 skip_bytes -= len(chunk)
243 self._pos = idx
244 self._eof = idx >= self._n
246 def seek(self, t: int | None = None, n: int | None = None) -> tuple["SeekResult", "EventArray", "TriggerArray | None"]:
247 from .common import SeekResult
248 if not self._is_initialized:
249 self.init()
250 axis, val = self._seek_axis(t, n)
252 # Close existing streams to prevent zipfile re-entry bugs on the same member
253 for fp in self._streams.values():
254 fp.close()
255 self._streams.clear()
257 if axis == "t":
258 ts = self._load_all_t()
259 idx = int(np.searchsorted(ts, val, side="left"))
260 else:
261 idx = val
262 idx = max(0, min(idx, self._n))
263 landed_ts = self._ts_at(idx)
264 self._seek_to_index(idx)
265 return SeekResult(ts=landed_ts if landed_ts is not None else val, index=idx, eof=self._eof), _EMPTY_EVENTS, None
267 def reset(self) -> None:
268 """Reset the reader to the beginning of the archive."""
269 if self._is_initialized:
270 self._open_streams()
271 self._pos = 0
272 self._eof = False
274 def tell(self) -> int:
275 """Current position, in events (npz has no meaningful byte offset)."""
276 return self._pos
278 def close(self) -> None:
279 """Close the member streams and the archive."""
280 for fp in self._streams.values():
281 fp.close()
282 self._streams = {}
283 if self._zf is not None:
284 self._zf.close()
285 self._zf = None
287class EventEncoder_Npz(EventEncoder):
288 """Encode events into an ``.npz`` archive with bounded memory.
290 Zip members can only be written one after another, while :meth:`write`
291 receives all four columns interleaved -- so each column is spooled to an
292 unlinked temporary file (raw bytes, no size limit from RAM) and the
293 archive is assembled on :meth:`close` by streaming every spool into its
294 ``.npy`` member.
296 Parameters
297 ----------
298 writable
299 Destination stream to write to.
300 width, height : int
301 Frame geometry stored in the archive.
302 dt : datetime, optional
303 Unused; npz stores no recording timestamp.
304 compressed : bool
305 Deflate the archive members (like ``np.savez_compressed``).
307 """
309 def __init__(self, writable: io.BufferedIOBase, width: int = 1280, height: int = 720,
310 dt: datetime | None = None, compressed: bool = False):
311 super().__init__(writable, width, height, dt)
312 self._compressed = compressed
313 self._spools: dict[str, IO[bytes]] = {}
314 self._closed = False
316 def init(self) -> None:
317 """Open one spool file per column (unlinked, cleaned up automatically)."""
318 if self._is_initialized:
319 return
320 self._spools = {name: tempfile.TemporaryFile() for name, _ in _COLUMNS}
321 self._is_initialized = True
323 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
324 """Append a chunk of events to the column spools.
326 Parameters
327 ----------
328 events : np.ndarray or EventArray
329 Array of events to write.
331 Returns
332 -------
333 int
334 Number of events written.
336 """
337 if not self._is_initialized:
338 self.init()
340 n = len(events)
341 for name, dtype in _COLUMNS:
342 col = np.ascontiguousarray(events[name], dtype=dtype)
343 self._spools[name].write(col.data)
344 self._n_written_events += n
345 return n
347 def flush(self) -> None:
348 """No-op: the archive can only be assembled once, on :meth:`close`."""
350 def close(self) -> None:
351 """Assemble the archive: stream each column spool into a ``.npy`` member."""
352 if self._closed:
353 return
354 self._closed = True
355 if not self._is_initialized:
356 self.init()
358 compression = zipfile.ZIP_DEFLATED if self._compressed else zipfile.ZIP_STORED
359 with zipfile.ZipFile(self._fd, "w", compression=compression, allowZip64=True) as zf:
360 for name, dtype in _COLUMNS:
361 spool = self._spools[name]
362 spool.flush()
363 spool.seek(0)
364 header = {
365 "descr": npy_format.dtype_to_descr(dtype),
366 "fortran_order": False,
367 "shape": (self._n_written_events,),
368 }
369 with zf.open(f"{name}.npy", "w", force_zip64=True) as dest:
370 npy_format.write_array_header_2_0(dest, header)
371 while True:
372 block = spool.read(1 << 22)
373 if not block:
374 break
375 dest.write(block)
376 spool.close()
377 self._spools = {}
379 for name, value in (("width", np.uint16(self._width)),
380 ("height", np.uint16(self._height))):
381 with zf.open(f"{name}.npy", "w") as dest:
382 npy_format.write_array(dest, np.asarray(value))
383 self._fd.flush()