Coverage for src/evutils/io/_csv.py: 92%
157 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +0000
1"""CSV file decoder and encoder."""
3import io
4import warnings
5from datetime import datetime
6from io import TextIOWrapper
7from pathlib import Path
8from typing import Any
10import numpy as np
11from ..types import Event_dtype, EventArray, TriggerArray
12from .common import EventDecoder, EventEncoder
15class EventDecoder_Csv(EventDecoder):
16 """A reader for CSV files with events.
18 Parameters
19 ----------
20 file
21 Path to the data file
22 order
23 Order of the columns in the CSV file, by default ['t', 'x', 'y', 'p']
24 chunk_size
25 Size of the chunk to read, by default 10_000
26 delimiter
27 Delimiter for the CSV file, by default ","
28 engine
29 Engine to use to read the CSV file, by default 'c'
31 Raises
32 ------
33 ValueError
34 If the order is not a list of 4 strings or if the order does not contain 't', 'x', 'y' and 'p'
36 """
38 def __init__(self, readable:io.BufferedReader, order:list[str]|None=None, chunk_size:int=1_000_000, delimiter:str=",", engine:str='c'):
39 super().__init__(readable, chunk_size)
42 if order is None:
43 # we infer the header from the file, or use a default header
44 pass
45 else:
46 # Validate the parameters
47 if len(order) != 4:
48 raise ValueError("Order must be a list of 4 strings")
49 if "t" not in order or "x" not in order or "y" not in order or "p" not in order:
50 raise ValueError("Order must contain 't', 'x', 'y' and 'p'")
52 self._order = order
53 self._delimiter = delimiter
54 self._engine = engine
58 def _check_header(self) -> None:
59 have_header = False
61 # Check first line to see if it is a header
62 # TODO: Need to deal with non-seekable files
63 first_line: str = self._fd.readline().decode('utf-8').strip()
65 # Check if the first line is a header
66 # If we find a header, it will take precendence over the order parameter
68 if "t" in first_line or "x" in first_line or "y" in first_line or "p" in first_line:
69 # Header found
70 have_header = True
72 if have_header:
73 # We expect a header:
74 if "t" in first_line and "x" in first_line and "y" in first_line and "p" in first_line:
75 # Header found
76 order = [c.strip() for c in first_line.split(self._delimiter)]
78 if self._order is not None and order != self._order:
79 warnings.warn(f"Header order {order} in file takes precedence over requested order {self._order}")
80 self._order = order
83 else:
84 raise ValueError(f"Header not found or invalid: {first_line}")
85 else:
86 # No header found, just seek to start of file
87 self._fd.seek(0)
89 # If we still don't have a header, we use a default one
90 if self._order is None:
91 self._order = ['t', 'x', 'y', 'p']
95 def init(self) -> None:
96 """Initialize the CSV reader.
98 Returns
99 -------
100 None
102 """
103 self._check_header()
104 assert self._order is not None # set by _check_header
106 # Build col_mapping: maps CSV index to out_array index
107 # out_arrays index: 0=t, 1=x, 2=y, 3=p
108 self._field_map = {'t': 0, 'x': 1, 'y': 2, 'p': 3}
109 self._col_mapping = [-1] * len(self._order)
110 for i, col in enumerate(self._order):
111 if col in self._field_map:
112 self._col_mapping[i] = self._field_map[col]
114 self._buffer = bytearray()
115 self._is_initialized = True
117 def read_chunk(self, delta_t_hint:int | None = None, n_events_hint:int | None = None) -> 'EventArray':
118 """Read a chunk of events from the CSV file."""
119 import ctypes
120 from . import _native_csv; from ._native_core import lib
122 assert self._is_initialized, "Reader is not initialized"
123 chunk_size = self._chunk_size
124 if n_events_hint is not None:
125 chunk_size = n_events_hint
127 t_arr = np.empty(chunk_size, dtype=np.int64)
128 x_arr = np.empty(chunk_size, dtype=np.uint16)
129 y_arr = np.empty(chunk_size, dtype=np.uint16)
130 p_arr = np.empty(chunk_size, dtype=np.uint8)
132 array_types = (ctypes.c_int * 4)(8, 2, 2, 1)
133 col_mapping = (ctypes.c_int * len(self._col_mapping))(*self._col_mapping)
134 delimiter = self._delimiter.encode('utf-8')[0]
136 events_parsed_total = 0
138 while events_parsed_total < chunk_size:
139 if not self._eof and len(self._buffer) < 1024 * 1024:
140 new_data = self._fd.read(4 * 1024 * 1024)
141 if not new_data:
142 self._eof = True
143 self._buffer.extend(b'\n') # Guarantee last line ends with newline
144 else:
145 self._buffer.extend(new_data)
147 if len(self._buffer) == 0:
148 break
150 bytes_consumed = ctypes.c_size_t(0)
151 events_parsed = ctypes.c_size_t(0)
153 cur_out_ptrs = (ctypes.c_void_p * 4)(
154 t_arr.ctypes.data + events_parsed_total * 8,
155 x_arr.ctypes.data + events_parsed_total * 2,
156 y_arr.ctypes.data + events_parsed_total * 2,
157 p_arr.ctypes.data + events_parsed_total * 1
158 )
160 c_buf = (ctypes.c_char * len(self._buffer)).from_buffer(self._buffer)
162 lib().evutils_read_csv(
163 c_buf, len(self._buffer), delimiter, cur_out_ptrs, array_types,
164 col_mapping, len(self._col_mapping), chunk_size - events_parsed_total,
165 ctypes.byref(bytes_consumed), ctypes.byref(events_parsed)
166 )
168 del c_buf # Release memory view so buffer can be resized
170 consumed = bytes_consumed.value
171 parsed = events_parsed.value
173 if consumed == 0:
174 if self._eof:
175 break
176 # No full line in the buffered window (a line longer than the
177 # refill threshold): pull more data regardless of the usual
178 # low-water mark, otherwise this loop would never progress.
179 new_data = self._fd.read(4 * 1024 * 1024)
180 if not new_data:
181 self._eof = True
182 self._buffer.extend(b'\n') # guarantee final-line newline
183 else:
184 self._buffer.extend(new_data)
185 continue
187 del self._buffer[:consumed]
188 events_parsed_total += parsed
190 if events_parsed_total < chunk_size:
191 self._eof = True
193 return EventArray(
194 t_arr[:events_parsed_total],
195 x_arr[:events_parsed_total],
196 y_arr[:events_parsed_total],
197 p_arr[:events_parsed_total]
198 )
200 def reset(self) -> None:
201 """Reset the CSV reader to the beginning of the file."""
202 assert self._fd is not None
203 self._fd.seek(0)
204 self._eof = False
205 if self._is_initialized:
206 self._is_initialized = False
207 self.init()
212class EventEncoder_Csv(EventEncoder):
213 """A writer for CSV files with events.
215 Parameters
216 ----------
217 file : str
218 Path to the data file
219 width : int, optional
220 Width of the frame, by default 1280 (not relevant for this formats)
221 height : int, optional
222 Height of the frame, by default 720 (not relevant for this formats)
223 sep : str, optional
224 Separator for the CSV file, by default ","
225 order : list, optional
226 Order of the columns in the CSV file, by default ['t', 'x', 'y', 'p']
227 header : bool, optional
228 If True, a header is written on the first line, by default True
230 Raises
231 ------
232 ValueError
233 If the order is not a list of 4 strings or if the order does not contain 't', 'x', 'y' and 'p'
235 """
237 def __init__(self, writable: io.BufferedWriter, width:int=1280, height:int=720, dt:datetime|None=None, sep:str=",", order:list[str]|None=None, header:bool=True):
238 super().__init__(writable, width, height, dt)
239 if order is None:
240 order = ['t', 'x', 'y', 'p']
243 if len(order) != 4:
244 raise ValueError("Order must be a list of 4 strings")
245 if "t" not in order or "x" not in order or "y" not in order or "p" not in order:
246 raise ValueError("Order must contain 't', 'x', 'y' and 'p'")
248 self._order = order
249 self._header = header
250 self._sep = sep
252 def init(self) -> None:
253 """Initialize the CSV writer.
255 Returns
256 -------
257 None
259 """
260 if self._header:
261 header = self._sep.join(self._order) + "\n"
262 self._fd.write(header.encode('utf-8'))
264 self._is_initialized = True
267 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
268 """Write events to the CSV file."""
269 import ctypes
270 from . import _native_csv; from ._native_core import lib
272 if not self._is_initialized:
273 self.init()
275 if isinstance(events, np.ndarray):
276 events = EventArray.from_aos(events)
278 chunk_size = len(events)
279 if chunk_size == 0:
280 return 0
282 t_arr = events.t
283 x_arr = events.x
284 y_arr = events.y
285 p_arr = events.p
287 in_ptrs_list = []
288 array_types_list = []
289 for col in self._order:
290 if col == 't':
291 in_ptrs_list.append(t_arr.ctypes.data)
292 array_types_list.append(8)
293 elif col == 'x':
294 in_ptrs_list.append(x_arr.ctypes.data)
295 array_types_list.append(2)
296 elif col == 'y':
297 in_ptrs_list.append(y_arr.ctypes.data)
298 array_types_list.append(2)
299 elif col == 'p':
300 in_ptrs_list.append(p_arr.ctypes.data)
301 array_types_list.append(1)
302 else:
303 in_ptrs_list.append(0)
304 array_types_list.append(0)
306 in_ptrs = (ctypes.c_void_p * len(self._order))(*in_ptrs_list)
307 array_types = (ctypes.c_int * len(self._order))(*array_types_list)
308 delimiter = self._sep.encode('utf-8')[0]
310 out_buffer_len = chunk_size * len(self._order) * 22
311 out_buffer = ctypes.create_string_buffer(out_buffer_len)
313 bytes_written = ctypes.c_size_t(0)
314 events_written = ctypes.c_size_t(0)
316 lib().evutils_write_csv(
317 in_ptrs, array_types, len(self._order), delimiter, chunk_size,
318 out_buffer, out_buffer_len, ctypes.byref(bytes_written), ctypes.byref(events_written)
319 )
321 self._fd.write(out_buffer.raw[:bytes_written.value])
322 return events_written.value