Coverage for src/evutils/io/_event_writer.py: 89%
87 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""Event writer module.
3Provides the `EventWriter` class for writing event data to various file formats.
4"""
6import io
7from datetime import datetime
8from pathlib import Path
10import numpy as np
12from . import encoders as ev_encoders
13from ._compression import is_compressed_path, open_compressed
15class EventWriter():
16 """Base class for writing events to different file formats.
18 Parameters
19 ----------
20 file
21 Path to the data file (or a writable binary stream). Compressed paths
22 (``.gz`` / ``.zst`` / ``.xz`` / ``.bz2``, e.g. ``foo.raw.zst``) are
23 compressed transparently; the inner extension selects the format.
24 Writing to a bare stream still requires an explicit ``file_encoder``.
25 width
26 Width of the frame. If None (default), taken from the first written
27 events' ``sensor_size`` metadata, falling back to 1280 (not relevant
28 for some formats).
29 height
30 Height of the frame. If None (default), taken from the first written
31 events' ``sensor_size`` metadata, falling back to 720 (not relevant
32 for some formats).
33 dt
34 Timestamp of the recording (default is the current time, but information is not saved in all formats)
35 file_encoder
36 File encoder to use, by default None (chosen from the file extension)
37 mode
38 File open mode, by default 'w+b'
39 **kwargs
40 Additional arguments for the file encoder
42 Examples
43 --------
44 >>> import numpy as np
45 >>> from evutils.types import EventArray
46 >>> # Create an EventArray
47 >>> events = EventArray(
48 ... t=np.array([0, 1000]),
49 ... x=np.array([0, 10]),
50 ... y=np.array([0, 10]),
51 ... p=np.array([1, 0])
52 ... )
53 >>> # Write the events to a raw file
54 >>> with EventWriter("events.raw") as writer: # doctest: +SKIP
55 ... writer.write(events) # doctest: +SKIP
57 """
59 def __init__(self, file: Path | str | io.BufferedIOBase, width:int|None=None, height:int|None=None, dt: datetime|None = None, file_encoder: ev_encoders.EventEncoder | None = None, mode: str = 'w+b', **kwargs):
60 self._mode = mode
61 self._file_name: Path | None = None
63 # Handle paths as input
64 if isinstance(file, str):
65 file = Path(file)
66 if isinstance(file, Path):
67 self._file_name = file
68 file = self._open_file(file)
69 else:
70 # A raw stream was passed - we need an explicit encoder.
71 if file_encoder is None:
72 raise ValueError("When using a binary stream as file, the file_encoder must be provided explicitly")
74 if isinstance(file, io.IOBase) and not file.writable():
75 raise IOError("File is not writable")
76 self._file: io.BufferedIOBase = file
78 # Encoder resolution is deferred: when width/height are not given
79 # explicitly we try to pick them up from the first written events'
80 # ``sensor_size`` metadata (falling back to 1280x720). An explicitly
81 # supplied encoder instance is used as-is.
82 self._explicit_width = width
83 self._explicit_height = height
84 self._raw_dt = dt # passed verbatim to the encoder
85 self._encoder_kwargs = kwargs
86 self._file_encoder = file_encoder # None => built lazily
88 self._width = width
89 self._height = height
90 self._n_written_events = 0
91 self._is_initialized = False
92 self._warned_triggers = False # warn once about unsupported triggers
93 self._dt = dt if dt is not None else datetime.now()
95 def _ensure_encoder(self, events: "EventArray | None" = None) -> None:
96 """Build the file encoder if it hasn't been resolved yet.
98 Dimensions are resolved in priority order: explicit ``width``/``height``
99 from the constructor, then the events' ``sensor_size`` metadata, then a
100 ``1280x720`` fallback. Formats that ignore geometry are unaffected.
101 """
102 if self._file_encoder is not None:
103 return
104 assert self._file_name is not None
105 w, h = self._explicit_width, self._explicit_height
106 if w is None or h is None:
107 ss = getattr(events, "sensor_size", None)
108 if ss is not None:
109 if w is None:
110 w = int(ss[0])
111 if h is None:
112 h = int(ss[1])
113 if w is None:
114 w = 1280
115 if h is None:
116 h = 720
117 self._width, self._height = w, h
118 encoder_cls = ev_encoders.get_file_writer(self._file_name)
119 self._file_encoder = encoder_cls(self._file, width=w, height=h,
120 dt=self._raw_dt, **self._encoder_kwargs)
122 def _open_file(self, file_name: Path) -> io.BufferedIOBase:
123 """Open the file for writing.
125 Parameters
126 ----------
127 file_name : Path
128 Path to the file to open.
130 Returns
131 -------
132 io.BufferedIOBase
133 The opened file object.
135 """
136 # Compressed paths (foo.raw.zst) auto-open through a compressing stream;
137 # the inner extension (resolved by get_file_writer) selects the encoder.
138 if is_compressed_path(file_name):
139 return open_compressed(file_name, "wb")
140 # default 'w+b' (not 'wb'): container encoders (HDF5) need the stream to be
141 # readable and seekable, and it costs nothing for the append-only ones.
142 return open(str(file_name), self._mode)
144 def init(self) -> None:
145 """Initialize the writer (e.g. open the file, write the header).
147 This method can be called explicitly, but it is also called automatically when the first event is written
148 """
149 self._ensure_encoder(None)
150 self._file_encoder.init()
151 self._is_initialized = True
153 def write(self, events: np.ndarray, triggers: np.ndarray | None = None) -> int:
154 """Write a buffer of events (and optionally external triggers) to the file.
156 Parameters
157 ----------
158 events
159 Buffer of events to write (structured array or EventArray)
160 triggers
161 Buffer of triggers to write (structured array or TriggerArray). Optional.
162 Written only by encoders that support trigger encoding (currently
163 AEDAT4, as separate trigger packets); all other formats ignore
164 them with a warning. To write a trigger-only batch, pass
165 ``EventArray.empty()`` as ``events``.
167 Returns
168 -------
169 int
170 Number of events written
172 Examples
173 --------
174 >>> from evutils.types import EventArray
175 >>> events = EventArray(t=[0, 100], x=[10, 11], y=[20, 21], p=[1, 0])
176 >>> writer = EventWriter("events.raw") # doctest: +SKIP
177 >>> num_written = writer.write(events) # doctest: +SKIP
178 >>> print(f"Wrote {num_written} events") # doctest: +SKIP
179 >>> writer.close() # doctest: +SKIP
181 """
182 # Resolve the encoder lazily so the first batch can supply sensor_size.
183 if self._file_encoder is None:
184 self._ensure_encoder(events)
186 if (triggers is not None and len(triggers) > 0
187 and not getattr(self._file_encoder, "SUPPORTS_WRITE_TRIGGERS", False)
188 and not self._warned_triggers):
189 import warnings
190 warnings.warn(
191 f"{self._file_encoder.__class__.__name__} does not support "
192 f"writing external triggers; they will NOT be stored.",
193 stacklevel=2,
194 )
195 self._warned_triggers = True
196 n_written = self._file_encoder.write(events, triggers=triggers)
197 self._n_written_events += n_written
198 return n_written
200 def flush(self) -> None:
201 """Flush the buffer to the file."""
202 if self._file_encoder is not None:
203 self._file_encoder.flush()
205 def __enter__(self) -> "EventWriter":
206 return self
208 def __repr__(self) -> str:
209 if self._is_initialized:
210 is_initialized_txt = f"Written {self._n_written_events} events"
211 else:
212 is_initialized_txt = "not initialized"
213 return f"{self.__class__.__name__}(file={self._file} - {is_initialized_txt}, {self._width}x{self._height})"
215 def __len__(self) -> int:
216 return self._n_written_events
218 def close(self) -> None:
219 """Close the writer and release the resources.
221 Finalizes the encoder first (container formats write their archive /
222 index here), then closes the underlying file.
223 """
224 # A writer that was opened but never written to still gets a valid
225 # (header-only) file, matching the pre-lazy-init behaviour.
226 self._ensure_encoder(None)
227 self._file_encoder.close()
228 if self._file_name is not None:
229 self._file.close()
231 def __exit__(self, exc_type: "type[BaseException] | None", exc_value: "BaseException | None", traceback: "types.TracebackType | None") -> None:
232 self.close()