Coverage for src/evutils/io/_event_writer.py: 78%
58 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +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
9from typing import Any
11import numpy as np
13from . import encoders as ev_encoders
16class EventWriter():
17 """Base class for writing events to different file formats.
19 Parameters
20 ----------
21 file
22 Path to the data file
23 width
24 Width of the frame, by default 1280 (not relevant for some formats)
25 height
26 Height of the frame, by default 720 (not relevant for some formats)
27 dt
28 Timestamp of the recording (default is the current time, but information is not saved in all formats)
29 file_encoder
30 File encoder to use, by default None (chosen from the file extension)
31 **kwargs
32 Additional arguments for the file encoder
35 Examples
36 --------
37 >>> import numpy as np
38 >>> from evutils.types import EventArray
39 >>> # Create an EventArray
40 >>> events = EventArray(
41 ... t=np.array([0, 1000]),
42 ... x=np.array([0, 10]),
43 ... y=np.array([0, 10]),
44 ... p=np.array([1, 0])
45 ... )
46 >>> # Write the events to a raw file
47 >>> with EventWriter("events.raw") as writer: # doctest: +SKIP
48 ... writer.write(events) # doctest: +SKIP
50 """
52 def __init__(self, file: Path | str | io.BufferedIOBase, width:int=1280, height:int=720, dt: datetime|None = None, file_encoder: ev_encoders.EventEncoder | None = None, **kwargs: Any):
54 self._file_name: Path | None = None
56 # Handle paths as input
57 if isinstance(file, str):
58 file = Path(file)
59 if isinstance(file, Path):
60 self._file_name = file
61 file = self._open_file(file)
62 else:
63 # A raw stream was passed - we need an explicit encoder.
64 if file_encoder is None:
65 raise ValueError("When using a binary stream as file, the file_encoder must be provided explicitly")
67 if isinstance(file, io.IOBase) and not file.writable():
68 raise IOError("File is not writable")
69 self._file: io.BufferedIOBase = file
71 # Resolve the encoder: explicit instance > heuristic from extension.
72 if file_encoder is None:
73 assert self._file_name is not None
74 encoder_cls = ev_encoders.get_file_writer(self._file_name)
75 self._file_encoder = encoder_cls(self._file, width=width, height=height, dt=dt, **kwargs)
76 else:
77 self._file_encoder = file_encoder
79 self._width = width
80 self._height = height
81 self._n_written_events = 0
82 self._is_initialized = False
83 self._dt = dt if dt is not None else datetime.now()
85 def _open_file(self, file_name: Path) -> io.BufferedIOBase:
86 """Open the file for writing.
88 Parameters
89 ----------
90 file_name : Path
91 Path to the file to open.
93 Returns
94 -------
95 io.BufferedIOBase
96 The opened file object.
98 """
99 # 'w+b' (not 'wb'): container encoders (HDF5) need the stream to be
100 # readable and seekable, and it costs nothing for the append-only ones.
101 return open(str(file_name), 'w+b')
103 def init(self) -> None:
104 """Initialize the writer (e.g. open the file, write the header).
106 This method can be called explicitly, but it is also called automatically when the first event is written
107 """
108 self._file_encoder.init()
109 self._is_initialized = True
111 def write(self, events: np.ndarray, triggers: np.ndarray | None = None) -> int:
112 """Write a buffer of events to the file.
114 Parameters
115 ----------
116 events
117 Buffer of events to write (structured array or EventArray)
118 triggers
119 Buffer of triggers to write (structured array or TriggerArray). Optional.
121 TODO: Add support for writing triggers to the file (currently only events are written)
123 Returns
124 -------
125 int
126 Number of events written
128 Examples
129 --------
130 >>> from evutils.types import EventArray
131 >>> events = EventArray(t=[0, 100], x=[10, 11], y=[20, 21], p=[1, 0])
132 >>> writer = EventWriter("events.raw") # doctest: +SKIP
133 >>> num_written = writer.write(events) # doctest: +SKIP
134 >>> print(f"Wrote {num_written} events") # doctest: +SKIP
135 >>> writer.close() # doctest: +SKIP
137 """
138 n_written = self._file_encoder.write(events)
139 self._n_written_events += n_written
140 return n_written
142 def write_external_triggers(self, triggers: np.ndarray) -> int:
143 """Write external triggers to the file.
145 Parameters
146 ----------
147 triggers : np.ndarray
148 Buffer of triggers to write (structured array or TriggerArray).
150 Returns
151 -------
152 int
153 Number of triggers written.
155 Examples
156 --------
157 >>> from evutils.types import TriggerArray
158 >>> triggers = TriggerArray(t=[1000, 2000], p=[1, 0], id=[0, 0])
159 >>> writer = EventWriter("events.raw") # doctest: +SKIP
160 >>> writer.write_external_triggers(triggers) # doctest: +SKIP
161 >>> writer.close() # doctest: +SKIP
163 """
164 from ..types import EventArray
165 self._file_encoder.write(EventArray.empty(), triggers=triggers)
166 return len(triggers)
168 def flush(self) -> None:
169 """Flush the buffer to the file."""
170 self._file_encoder.flush()
172 def __enter__(self) -> "EventWriter":
173 return self
175 def __repr__(self) -> str:
176 if self._is_initialized:
177 is_initialized_txt = f"Written {self._n_written_events} events"
178 else:
179 is_initialized_txt = "not initialized"
180 return f"{self.__class__.__name__}(file={self._file} - {is_initialized_txt}, {self._width}x{self._height})"
182 def __len__(self) -> int:
183 return self._n_written_events
185 def close(self) -> None:
186 """Close the writer and release the resources.
188 Finalizes the encoder first (container formats write their archive /
189 index here), then closes the underlying file.
190 """
191 self._file_encoder.close()
192 self._file.close()
194 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
195 self.close()