Coverage for src/evutils/io/common.py: 76%
109 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +0000
1"""Common interface for event decoders and encoders.
3Defines the abstract base classes `EventDecoder` and `EventEncoder`.
4"""
6import io
7from abc import ABC, abstractmethod
8from datetime import datetime
9from typing import Any, Optional, TYPE_CHECKING
11if TYPE_CHECKING:
12 from ..types import EventArray, TriggerArray
14import numpy as np
19class EventDecoder(ABC):
20 """ABC for reading chunks of events from a IO source object.
22 Parameters
23 ----------
24 readable
25 source to read events from
26 chunk_size
27 Size of the chunk to read
29 Raises
30 ------
31 NotImplementedError
32 If the method is not implemented in the subclass
34 """
36 SUPPORTS_EXT_TRIGGERS = False
38 def __init__(self, source: Any, chunk_size: int = 10000, read_external_triggers: bool = False):
39 """Initialize the decoder.
41 Parameters
42 ----------
43 source : ByteSource
44 Source to read events from.
45 chunk_size : int, optional
46 Size of the chunk to read, by default 10000.
47 read_external_triggers : bool, optional
48 Whether to read external triggers, by default False.
50 """
51 # `source` is a ByteSource (see io/_source.py). `fd` is kept as a legacy
52 # alias for older decoders that still reference it.
53 self._source = source
54 self._fd = source
56 self._is_initialized = False
58 self._chunk_size = chunk_size
60 self._eof = False
62 self._width: int | None = None
63 self._height: int | None = None
65 self.read_external_triggers = read_external_triggers
66 if self.read_external_triggers and not self.SUPPORTS_EXT_TRIGGERS:
67 import warnings
68 warnings.warn(f"{self.__class__.__name__} does not support reading external triggers.")
70 @abstractmethod
71 def init(self) -> None:
72 """Initialize the file for reading."""
73 raise NotImplementedError
75 @abstractmethod
76 def read_chunk(self, delta_t_hint:int | None = None, n_events_hint:int | None = None) -> 'EventArray | tuple[EventArray, TriggerArray]':
77 """Read a chunk of events.
79 Parameters
80 ----------
81 delta_t_hint : int, optional
82 If not None, can be used to provide a hit about the delta_t window to be read
83 n_events_hint : int, optional
84 If not None, can be used to provide a hit about the n_events to be read
86 Returns
87 -------
88 EventArray or tuple of (EventArray, TriggerArray)
89 Chunk of events read from the source. Optionally returns triggers if `read_external_triggers` is True.
91 """
92 raise NotImplementedError
94 @abstractmethod
95 def reset(self) -> None:
96 """Reset the file pointer to the beginning of the file."""
97 raise NotImplementedError
100 def read_all(self) -> 'EventArray | tuple[EventArray, TriggerArray]':
101 """Decode and return every remaining event at once.
103 The default implementation drains :meth:`read_chunk` and concatenates the
104 chunks. SoA-native decoders (EVT/DAT/AER) override this with a
105 single-buffer decode that avoids the per-chunk copy entirely.
107 Returns
108 -------
109 EventArray
110 All remaining events.
112 """
113 from ..types import EventArray
115 if not self._is_initialized:
116 self.init()
118 # read_chunk may return a view that is invalidated by the next call, so
119 # copy each chunk before pulling the next one.
120 chunks = []
121 trigger_chunks = []
122 while True:
123 _chunk = self.read_chunk()
124 if self.read_external_triggers:
125 if isinstance(_chunk, tuple):
126 ev_chunk, tr_chunk = _chunk
127 else:
128 ev_chunk = _chunk
129 from ..types import TriggerArray
130 tr_chunk = TriggerArray.empty()
131 else:
132 ev_chunk = _chunk # type: ignore
133 if len(ev_chunk) == 0:
134 # The final chunk can carry triggers with no events (trigger
135 # packets after the last CD event): keep them before stopping.
136 if self.read_external_triggers and len(tr_chunk) > 0:
137 trigger_chunks.append(tr_chunk.copy())
138 break
139 chunks.append(ev_chunk.copy())
140 if self.read_external_triggers:
141 trigger_chunks.append(tr_chunk.copy())
143 if not chunks:
144 res_ev = EventArray.empty()
145 elif len(chunks) == 1:
146 res_ev = chunks[0]
147 else:
148 res_ev = EventArray(
149 np.concatenate([c.t for c in chunks]),
150 np.concatenate([c.x for c in chunks]),
151 np.concatenate([c.y for c in chunks]),
152 np.concatenate([c.p for c in chunks]),
153 )
155 if self.read_external_triggers:
156 if not trigger_chunks:
157 from ..types import TriggerArray
158 res_tr = TriggerArray.empty()
159 elif len(trigger_chunks) == 1:
160 res_tr = trigger_chunks[0]
161 else:
162 from ..types import TriggerArray
163 res_tr = TriggerArray(
164 np.concatenate([c.t for c in trigger_chunks]),
165 np.concatenate([c.p for c in trigger_chunks]),
166 np.concatenate([c.id for c in trigger_chunks]),
167 )
168 return res_ev, res_tr
170 return res_ev
172 def close(self) -> None:
173 """Release any resources held by the decoder (e.g. buffer views).
175 The owning source is closed separately by the EventReader.
176 """
177 pass
179 def tell(self) -> int:
180 """Get the current position in the file.
182 Returns
183 -------
184 int
185 The current position in the file
187 """
188 return int(self._fd.tell())
191 def set_chunk_size(self, chunk_size:int) -> None:
192 """Set the chunk size.
194 Parameters
195 ----------
196 chunk_size
197 Size of the chunk to read
199 """
200 self._chunk_size = chunk_size
202 def shape(self) -> tuple[int|None, int|None]:
203 """Get the shape of the frame (width, height).
205 Returns
206 -------
207 tuple[int|None, int|None]
208 The shape of the frame (width, height), or (None, None) if the shape is not known
210 """
211 return self._width, self._height
215 def __repr__(self) -> str:
216 if self._is_initialized:
217 is_initialized_txt = "initialized"
218 else:
219 is_initialized_txt = "not initialized"
220 return f"{self.__class__} - {is_initialized_txt})"
222 def is_eof(self) -> bool:
223 """Check if the end of the file has been reached.
225 Returns
226 -------
227 bool
228 True if the end of the file has been reached
230 """
231 return self._eof
236class EventEncoder(ABC):
237 """ABC for writing chunks of events to a io object.
239 Parameters
240 ----------
241 writable
242 Destination for writing events
243 width : int, optional
244 Width of the frame, by default 1280 (not relevant for some formats)
245 height : int, optional
246 Height of the frame, by default 720 (not relevant for some formats)
247 dt : datetime, optional
248 Timestamp of the recording (default is the current time, but information is not saved in all formats)
251 Raises
252 ------
253 NotImplementedError
254 If the method is not implemented in the subclass
256 """
258 def __init__(self, writable: io.BufferedIOBase, width:int = 1280, height:int = 720, dt:Optional[datetime]=None ):
259 """Initialize the encoder.
261 Parameters
262 ----------
263 writable : io.BufferedIOBase
264 Destination for writing events.
265 width : int, optional
266 Width of the frame, by default 1280.
267 height : int, optional
268 Height of the frame, by default 720.
269 dt : datetime, optional
270 Timestamp of the recording, by default current time.
272 """
273 self._fd = writable
275 self._width = width
276 self._height = height
279 self._n_written_events = 0
280 self._is_initialized = False
282 if dt is None:
283 self._dt = datetime.now()
284 else:
285 self._dt = dt
287 @abstractmethod
288 def init(self) -> None:
289 """Initialize the file for writing."""
290 raise NotImplementedError
292 @abstractmethod
293 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
294 """Write a chunk of events."""
295 raise NotImplementedError
297 def __len__(self) -> int:
298 return self._n_written_events
301 def __enter__(self) -> "EventEncoder":
302 return self
305 def __repr__(self) -> str:
306 if self._is_initialized:
307 is_initialized_txt = f"Written {self._n_written_events} events"
308 else:
309 is_initialized_txt = "not initialized"
310 return f"{self.__class__.__name__} - {is_initialized_txt}, {self._width}x{self._height})"
312 def flush(self) -> None:
313 """Flush any buffered data to the underlying stream."""
314 self._fd.flush()
316 def close(self) -> None:
317 """Finalize the encoder.
319 Container formats (NPZ, HDF5) override this to write the archive /
320 index before the underlying file is closed. The owning stream itself is
321 closed by the :class:`~evutils.io.EventWriter`, not here.
322 """
323 self.flush()