Coverage for src/evutils/io/_npz.py: 93%
165 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +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, Any
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)))
41def _read_npy_header(fp: IO[bytes]) -> tuple[tuple[int, ...], np.dtype]:
42 """Read the ``.npy`` magic + header from a stream, returning (shape, dtype).
44 Leaves ``fp`` positioned at the first data byte. Fortran-ordered arrays are
45 rejected (event columns are 1-D, written C-ordered).
46 """
47 version = npy_format.read_magic(fp)
48 read_header = {
49 (1, 0): npy_format.read_array_header_1_0,
50 (2, 0): npy_format.read_array_header_2_0,
51 }.get(version)
52 if read_header is None:
53 raise ValueError(f"unsupported .npy format version {version}")
54 shape, fortran, dtype = read_header(fp)
55 if fortran:
56 raise ValueError("Fortran-ordered .npy members are not supported")
57 return shape, dtype
60def _read_exact(fp: IO[bytes], nbytes: int) -> bytearray:
61 """Read exactly ``nbytes`` from a (possibly decompressing) stream.
63 Returns a writable buffer so the numpy views over it are mutable.
64 """
65 out = bytearray(nbytes)
66 view = memoryview(out)
67 got = 0
68 while got < nbytes:
69 n = fp.readinto(view[got:]) # type: ignore[attr-defined]
70 if not n:
71 raise EOFError(f"truncated .npy member: expected {nbytes} bytes, got {got}")
72 got += n
73 return out
76class EventDecoder_Npz(EventDecoder):
77 """Decode events from an ``.npz`` archive, streaming chunk by chunk.
79 The archive members are read through zip streams: only ``chunk_size``
80 events are held in memory at a time, so arbitrarily large recordings can
81 be iterated. Accepts either the four column members ``t/x/y/p`` or a
82 single structured ``events`` member.
84 Parameters
85 ----------
86 source
87 Byte source to read from (must be seekable, as required by the zip
88 format).
89 chunk_size
90 Number of events returned per :meth:`read_chunk` call.
92 """
94 def __init__(self, source: ByteSource, chunk_size: int = 1_000_000):
95 super().__init__(source, chunk_size)
96 self._zf: zipfile.ZipFile | None = None
97 self._streams: dict[str, IO[bytes]] = {}
98 self._aos_dtype: np.dtype | None = None # set when reading an 'events' member
99 self._n = 0
100 self._pos = 0
102 def _open_streams(self) -> None:
103 """(Re)open the member streams and position them past the npy headers."""
104 assert self._zf is not None
105 for fp in self._streams.values():
106 fp.close()
107 self._streams = {}
109 names = set(self._zf.namelist())
110 if {"t.npy", "x.npy", "y.npy", "p.npy"} <= names:
111 n = None
112 for name, _ in _COLUMNS:
113 fp = self._zf.open(f"{name}.npy")
114 shape, dtype = _read_npy_header(fp)
115 if len(shape) != 1:
116 raise ValueError(f"member {name}.npy is not 1-D: shape {shape}")
117 if n is None:
118 n = shape[0]
119 elif shape[0] != n:
120 raise ValueError("event columns have mismatched lengths")
121 self._streams[name] = fp
122 self._n = n or 0
123 self._aos_dtype = None
124 elif "events.npy" in names:
125 fp = self._zf.open("events.npy")
126 shape, dtype = _read_npy_header(fp)
127 if dtype.names is None or not {"t", "x", "y", "p"} <= set(dtype.names):
128 raise ValueError("'events' member must be a structured array with t/x/y/p fields")
129 self._streams["events"] = fp
130 self._aos_dtype = dtype
131 self._n = shape[0]
132 else:
133 raise ValueError(
134 f"NPZ archive does not contain event data: expected members "
135 f"'t/x/y/p' or 'events', found {sorted(names)}"
136 )
138 def init(self) -> None:
139 """Open the archive and locate the event members."""
140 if self._is_initialized:
141 return
143 f: Any = self._source if self._source.seekable() else io.BytesIO(self._source.read(-1))
144 self._zf = zipfile.ZipFile(f)
146 names = set(self._zf.namelist())
147 for attr, member in (("_width", "width.npy"), ("_height", "height.npy")):
148 if member in names:
149 with self._zf.open(member) as fp:
150 setattr(self, attr, int(npy_format.read_array(fp).item()))
152 self._open_streams()
153 self._pos = 0
154 self._is_initialized = True
156 def _read_n(self, n: int) -> EventArray:
157 """Stream the next ``n`` events out of the member streams."""
158 if self._aos_dtype is not None:
159 fp = self._streams["events"]
160 buf = _read_exact(fp, n * self._aos_dtype.itemsize)
161 return EventArray.from_aos(np.frombuffer(buf, dtype=self._aos_dtype))
162 cols = {}
163 for name, dtype in _COLUMNS:
164 # The member's own dtype was validated against 1-D at open; event
165 # columns are cast to the canonical dtypes by EventArray.
166 buf = _read_exact(self._streams[name], n * dtype.itemsize)
167 cols[name] = np.frombuffer(buf, dtype=dtype)
168 return EventArray(cols["t"], cols["x"], cols["y"], cols["p"])
170 def read_chunk(self, delta_t_hint: int | None = None,
171 n_events_hint: int | None = None) -> EventArray:
172 if not self._is_initialized:
173 self.init()
175 if self._pos >= self._n:
176 self._eof = True
177 return _EMPTY_EVENTS
179 n = min(self._chunk_size, self._n - self._pos)
180 chunk = self._read_n(n)
181 self._pos += n
182 if self._pos >= self._n:
183 self._eof = True
184 return chunk
186 def reset(self) -> None:
187 """Reset the reader to the beginning of the archive."""
188 if self._is_initialized:
189 self._open_streams()
190 self._pos = 0
191 self._eof = False
193 def tell(self) -> int:
194 """Current position, in events (npz has no meaningful byte offset)."""
195 return self._pos
197 def close(self) -> None:
198 """Close the member streams and the archive."""
199 for fp in self._streams.values():
200 fp.close()
201 self._streams = {}
202 if self._zf is not None:
203 self._zf.close()
204 self._zf = None
207class EventEncoder_Npz(EventEncoder):
208 """Encode events into an ``.npz`` archive with bounded memory.
210 Zip members can only be written one after another, while :meth:`write`
211 receives all four columns interleaved -- so each column is spooled to an
212 unlinked temporary file (raw bytes, no size limit from RAM) and the
213 archive is assembled on :meth:`close` by streaming every spool into its
214 ``.npy`` member.
216 Parameters
217 ----------
218 writable
219 Destination stream to write to.
220 width, height : int
221 Frame geometry stored in the archive.
222 dt : datetime, optional
223 Unused; npz stores no recording timestamp.
224 compressed : bool
225 Deflate the archive members (like ``np.savez_compressed``).
227 """
229 def __init__(self, writable: io.BufferedIOBase, width: int = 1280, height: int = 720,
230 dt: datetime | None = None, compressed: bool = False):
231 super().__init__(writable, width, height, dt)
232 self._compressed = compressed
233 self._spools: dict[str, IO[bytes]] = {}
234 self._closed = False
236 def init(self) -> None:
237 """Open one spool file per column (unlinked, cleaned up automatically)."""
238 if self._is_initialized:
239 return
240 self._spools = {name: tempfile.TemporaryFile() for name, _ in _COLUMNS}
241 self._is_initialized = True
243 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
244 """Append a chunk of events to the column spools.
246 Parameters
247 ----------
248 events : np.ndarray or EventArray
249 Array of events to write.
251 Returns
252 -------
253 int
254 Number of events written.
256 """
257 if not self._is_initialized:
258 self.init()
260 n = len(events)
261 for name, dtype in _COLUMNS:
262 col = np.ascontiguousarray(events[name], dtype=dtype)
263 self._spools[name].write(col.data)
264 self._n_written_events += n
265 return n
267 def flush(self) -> None:
268 """No-op: the archive can only be assembled once, on :meth:`close`."""
270 def close(self) -> None:
271 """Assemble the archive: stream each column spool into a ``.npy`` member."""
272 if self._closed:
273 return
274 self._closed = True
275 if not self._is_initialized:
276 self.init()
278 compression = zipfile.ZIP_DEFLATED if self._compressed else zipfile.ZIP_STORED
279 with zipfile.ZipFile(self._fd, "w", compression=compression, allowZip64=True) as zf:
280 for name, dtype in _COLUMNS:
281 spool = self._spools[name]
282 spool.flush()
283 spool.seek(0)
284 header = {
285 "descr": npy_format.dtype_to_descr(dtype),
286 "fortran_order": False,
287 "shape": (self._n_written_events,),
288 }
289 with zf.open(f"{name}.npy", "w", force_zip64=True) as dest:
290 npy_format.write_array_header_2_0(dest, header)
291 while True:
292 block = spool.read(1 << 22)
293 if not block:
294 break
295 dest.write(block)
296 spool.close()
297 self._spools = {}
299 for name, value in (("width", np.uint16(self._width)),
300 ("height", np.uint16(self._height))):
301 with zf.open(f"{name}.npy", "w") as dest:
302 npy_format.write_array(dest, np.asarray(value))
303 self._fd.flush()