Coverage for src/evutils/io/_csv.py: 90%
251 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""CSV file decoder and encoder."""
3import io
4import warnings
5from datetime import datetime
6from io import TextIOWrapper
7from pathlib import Path
9import numpy as np
10from ..types import Event_dtype, EventArray, TriggerArray
11from .common import EventDecoder, EventEncoder
13_EMPTY_EVENTS = EventArray.empty()
15class EventDecoder_Csv(EventDecoder):
16 """A reader for CSV files with events.
18 Parameters
19 ----------
20 source
21 Byte source to read from.
22 order
23 Order of the columns in the CSV file, by default ['t', 'x', 'y', 'p']
24 chunk_size
25 Maximum number of events per chunk, by default 1_000_000
26 delimiter
27 Delimiter for the CSV file, by default ","
29 Raises
30 ------
31 ValueError
32 If the order is not a list of 4 strings or if the order does not contain 't', 'x', 'y' and 'p'
34 """
36 #: read_chunk parses into fresh, independent arrays bounded by n_events_hint,
37 #: so EventReader can hand them out directly (skipping the staging
38 #: accumulator). CSV decode is text-parse-bound, so the gain is small.
39 _independent_windows = True
41 #: Seekable via byte-offset binary search with newline resync (time) or a
42 #: newline count (event index). Requires a seekable source.
43 SUPPORTS_SEEK = True
45 def __init__(self, source: io.BufferedReader, order: list[str] | None = None,
46 chunk_size: int = 1_000_000, delimiter: str = ","):
47 super().__init__(source, chunk_size)
49 if order is None:
50 # we infer the header from the file, or use a default header
51 pass
52 else:
53 # Validate the parameters
54 if len(order) != 4:
55 raise ValueError("Order must be a list of 4 strings")
56 if "t" not in order or "x" not in order or "y" not in order or "p" not in order:
57 raise ValueError("Order must contain 't', 'x', 'y' and 'p'")
59 self._order = order
60 self._delimiter = delimiter
62 def _check_header(self) -> None:
63 have_header = False
65 # Check the first line to see if it is a header, then rewind. A
66 # non-seekable source has already been slurped into a seekable BytesIO
67 # by init() (self._fd), so readline()/seek(0) here are always valid.
68 first_line: str = self._fd.readline().decode('utf-8').strip()
70 # Check if the first line is a header
71 # If we find a header, it will take precendence over the order parameter
73 cols = [c.strip() for c in first_line.split(self._delimiter)]
74 if "t" in cols or "x" in cols or "y" in cols or "p" in cols:
75 # Header found
76 have_header = True
78 if have_header:
79 # We expect a header:
80 if "t" in cols and "x" in cols and "y" in cols and "p" in cols:
81 # Header found
82 order = cols
84 if self._order is not None and order != self._order:
85 warnings.warn(f"Header order {order} in file takes precedence over requested order {self._order}")
86 self._order = order
88 else:
89 raise ValueError(f"Header not found or invalid: {first_line}")
90 else:
91 # No header found, just seek to start of file
92 self._fd.seek(0)
94 # If we still don't have a header, we use a default one
95 if self._order is None:
96 self._order = ['t', 'x', 'y', 'p']
98 def init(self) -> None:
99 """Initialize the CSV reader.
101 Returns
102 -------
103 None
105 """
106 # Header handling and byte-offset seeking need a seekable source. A
107 # non-seekable source (pipe, compressed stream) is slurped into an
108 # in-memory buffer first (mirrors the NPZ decoder).
109 if not self._source.seekable():
110 self._fd = io.BytesIO(self._source.read(-1))
112 self._check_header()
113 assert self._order is not None # set by _check_header
115 # Build col_mapping: maps CSV index to out_array index
116 # out_arrays index: 0=t, 1=x, 2=y, 3=p
117 self._field_map = {'t': 0, 'x': 1, 'y': 2, 'p': 3}
118 self._col_mapping = [-1] * len(self._order)
119 for i, col in enumerate(self._order):
120 if col in self._field_map:
121 self._col_mapping[i] = self._field_map[col]
123 # Byte offset of the first event line (past any header), for seeking.
124 self._data_start = int(self._fd.tell())
126 self._buffer = bytearray()
127 self._is_initialized = True
129 def read_chunk(self, delta_t_hint:int | None = None, n_events_hint:int | None = None) -> 'EventArray':
130 """Read a chunk of events from the CSV file."""
131 import ctypes
132 from . import _native_csv; from ._native_core import lib
134 assert self._is_initialized, "Reader is not initialized"
135 chunk_size = self._chunk_size
136 if n_events_hint is not None:
137 chunk_size = n_events_hint
139 t_arr = np.zeros(chunk_size, dtype=np.int64)
140 x_arr = np.zeros(chunk_size, dtype=np.uint16)
141 y_arr = np.zeros(chunk_size, dtype=np.uint16)
142 p_arr = np.zeros(chunk_size, dtype=np.uint8)
144 array_types = (ctypes.c_int * 4)(8, 2, 2, 1)
145 col_mapping = (ctypes.c_int * len(self._col_mapping))(*self._col_mapping)
146 delimiter = self._delimiter.encode('utf-8')[0]
148 events_parsed_total = 0
149 malformed_total = 0
151 while events_parsed_total < chunk_size:
152 if not self._eof and len(self._buffer) < 1024 * 1024:
153 new_data = self._fd.read(4 * 1024 * 1024)
154 if not new_data:
155 self._eof = True
156 self._buffer.extend(b'\n') # Guarantee last line ends with newline
157 else:
158 self._buffer.extend(new_data)
160 if len(self._buffer) == 0:
161 break
163 events_parsed = ctypes.c_size_t(0)
164 malformed = ctypes.c_size_t(0)
166 cur_out_ptrs = (ctypes.c_void_p * 4)(
167 t_arr.ctypes.data + events_parsed_total * 8,
168 x_arr.ctypes.data + events_parsed_total * 2,
169 y_arr.ctypes.data + events_parsed_total * 2,
170 p_arr.ctypes.data + events_parsed_total * 1
171 )
173 c_buf = (ctypes.c_char * len(self._buffer)).from_buffer(self._buffer)
175 res = lib().evutils_read_csv(
176 c_buf, len(self._buffer), delimiter, cur_out_ptrs, array_types,
177 col_mapping, len(self._col_mapping), chunk_size - events_parsed_total,
178 ctypes.byref(events_parsed), ctypes.byref(malformed)
179 )
181 consumed = (res.current - ctypes.addressof(c_buf)) if res.current is not None else 0
183 del c_buf # Release memory view so buffer can be resized
185 parsed = events_parsed.value
186 malformed_total += malformed.value
188 if consumed == 0:
189 if self._eof:
190 break
191 # No full line in the buffered window (a line longer than the
192 # refill threshold): pull more data regardless of the usual
193 # low-water mark, otherwise this loop would never progress.
194 new_data = self._fd.read(4 * 1024 * 1024)
195 if not new_data:
196 self._eof = True
197 self._buffer.extend(b'\n') # guarantee final-line newline
198 else:
199 self._buffer.extend(new_data)
200 continue
202 del self._buffer[:consumed]
203 events_parsed_total += parsed
205 if events_parsed_total < chunk_size:
206 self._eof = True
208 # Rows missing a mapped column were zero-filled (lenient default). Surface
209 # them the same way the binary parsers surface malformed packets: raise in
210 # strict mode, else warn once for the chunk.
211 if malformed_total:
212 from ._native_core import _handle_parse_warning
213 _handle_parse_warning(f"in {malformed_total} CSV row(s)", self._strict)
215 return EventArray(
216 t_arr[:events_parsed_total],
217 x_arr[:events_parsed_total],
218 y_arr[:events_parsed_total],
219 p_arr[:events_parsed_total]
220 )
222 # ------------------------------------------------------------------ #
223 # Seeking (byte-offset binary search + newline resync)
224 # ------------------------------------------------------------------ #
225 def _file_size(self) -> int:
226 cur = self._fd.tell()
227 end = self._fd.seek(0, io.SEEK_END)
228 self._fd.seek(cur)
229 return end
231 def _line_start_at_or_after(self, pos: int) -> int:
232 """Byte offset of the first line whose start is ``>= pos``."""
233 if pos <= self._data_start:
234 return self._data_start
235 self._fd.seek(pos - 1) # find the newline at/after pos-1; line starts after it
236 acc = 0
237 while True:
238 chunk = self._fd.read(65536)
239 if not chunk:
240 return pos - 1 + acc # EOF, no newline
241 i = chunk.find(b"\n")
242 if i >= 0:
243 return pos - 1 + acc + i + 1
244 acc += len(chunk)
246 def _parse_line_t(self, pos: int) -> int | None:
247 """Parse the ``t`` column of the line starting at ``pos`` (or ``None``)."""
248 assert self._order is not None
249 self._fd.seek(pos)
250 line = b""
251 while True:
252 chunk = self._fd.read(65536)
253 if not chunk:
254 break
255 nl = chunk.find(b"\n")
256 if nl >= 0:
257 line += chunk[:nl]
258 break
259 line += chunk
260 if not line.strip():
261 return None
262 parts = line.split(self._delimiter.encode("utf-8"))
263 ti = self._order.index("t")
264 try:
265 return int(parts[ti].strip())
266 except (ValueError, IndexError):
267 return None
269 def _seek_line_index(self, n: int) -> int:
270 """Byte offset of event line ``n`` (0-based), by counting newlines."""
271 self._fd.seek(self._data_start)
272 pos = self._data_start
273 remaining = n
274 while remaining > 0:
275 chunk = self._fd.read(1 << 20)
276 if not chunk:
277 break
278 cnt = chunk.count(b"\n")
279 if cnt < remaining:
280 remaining -= cnt
281 pos += len(chunk)
282 else:
283 idx = -1
284 for _ in range(remaining):
285 idx = chunk.find(b"\n", idx + 1)
286 pos += idx + 1
287 remaining = 0
288 return pos
290 def seek(self, t: int | None = None, n: int | None = None) -> tuple["SeekResult", "EventArray", "TriggerArray | None"]:
291 """Seek to an absolute timestamp (µs) or event index. See base class."""
292 from .common import SeekResult
293 if not self._is_initialized:
294 self.init()
295 axis, val = self._seek_axis(t, n)
296 size = self._file_size()
298 if axis == "n":
299 pos = self._seek_line_index(val)
300 idx = val
301 else:
302 lo, hi = self._data_start, size
303 while lo < hi:
304 mid = (lo + hi) // 2
305 ls = self._line_start_at_or_after(mid)
306 tv = self._parse_line_t(ls) if ls < size else None
307 if tv is not None and tv >= val:
308 hi = mid
309 else:
310 lo = mid + 1
311 pos = self._line_start_at_or_after(lo)
312 idx = -1
314 self._fd.seek(pos)
315 self._buffer = bytearray()
316 self._eof = False
317 landed_ts = self._parse_line_t(pos) if pos < size else None
318 self._fd.seek(pos) # _parse_line_t moved the cursor
320 # If pos >= size, we are at EOF.
321 if pos >= size:
322 self._eof = True
324 return SeekResult(ts=landed_ts if landed_ts is not None else val, index=idx, eof=self._eof), _EMPTY_EVENTS, None
326 def reset(self) -> None:
327 """Reset the CSV reader to the beginning of the file."""
328 assert self._fd is not None
329 self._fd.seek(0)
330 self._eof = False
331 if self._is_initialized:
332 self._is_initialized = False
333 self.init()
335class EventEncoder_Csv(EventEncoder):
336 """A writer for CSV files with events.
338 Parameters
339 ----------
340 file : str
341 Path to the data file
342 width : int, optional
343 Width of the frame, by default 1280 (not relevant for this formats)
344 height : int, optional
345 Height of the frame, by default 720 (not relevant for this formats)
346 sep : str, optional
347 Separator for the CSV file, by default ","
348 order : list, optional
349 Order of the columns in the CSV file, by default ['t', 'x', 'y', 'p']
350 header : bool, optional
351 If True, a header is written on the first line, by default True
353 Raises
354 ------
355 ValueError
356 If the order is not a list of 4 strings or if the order does not contain 't', 'x', 'y' and 'p'
358 """
360 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):
361 super().__init__(writable, width, height, dt)
362 if order is None:
363 order = ['t', 'x', 'y', 'p']
365 if len(order) != 4:
366 raise ValueError("Order must be a list of 4 strings")
367 if "t" not in order or "x" not in order or "y" not in order or "p" not in order:
368 raise ValueError("Order must contain 't', 'x', 'y' and 'p'")
370 self._order = order
371 self._header = header
372 self._sep = sep
374 def init(self) -> None:
375 """Initialize the CSV writer.
377 Returns
378 -------
379 None
381 """
382 if self._header:
383 header = self._sep.join(self._order) + "\n"
384 self._fd.write(header.encode('utf-8'))
386 self._is_initialized = True
388 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:
389 """Write events to the CSV file."""
390 import ctypes
391 from . import _native_csv; from ._native_core import lib
393 if not self._is_initialized:
394 self.init()
396 if isinstance(events, np.ndarray):
397 events = EventArray.from_aos(events)
399 chunk_size = len(events)
400 if chunk_size == 0:
401 return 0
403 t_arr = events.t
404 x_arr = events.x
405 y_arr = events.y
406 p_arr = events.p
408 in_ptrs_list = []
409 array_types_list = []
410 for col in self._order:
411 if col == 't':
412 in_ptrs_list.append(t_arr.ctypes.data)
413 array_types_list.append(8)
414 elif col == 'x':
415 in_ptrs_list.append(x_arr.ctypes.data)
416 array_types_list.append(2)
417 elif col == 'y':
418 in_ptrs_list.append(y_arr.ctypes.data)
419 array_types_list.append(2)
420 elif col == 'p':
421 in_ptrs_list.append(p_arr.ctypes.data)
422 array_types_list.append(1)
423 else:
424 in_ptrs_list.append(0)
425 array_types_list.append(0)
427 in_ptrs = (ctypes.c_void_p * len(self._order))(*in_ptrs_list)
428 array_types = (ctypes.c_int * len(self._order))(*array_types_list)
429 delimiter = self._sep.encode('utf-8')[0]
431 out_buffer_len = chunk_size * len(self._order) * 22
432 out_buffer = ctypes.create_string_buffer(out_buffer_len)
434 bytes_written = ctypes.c_size_t(0)
435 events_written = ctypes.c_size_t(0)
437 lib().evutils_write_csv(
438 in_ptrs, array_types, len(self._order), delimiter, chunk_size,
439 out_buffer, out_buffer_len, ctypes.byref(bytes_written), ctypes.byref(events_written)
440 )
442 self._fd.write(out_buffer.raw[:bytes_written.value])
443 return events_written.value