Coverage for src/evutils/io/stream.py: 100%

1from typing import Iterator 

2 

3from evutils.io._source import make_source 

4from evutils.io.decoders import resolve_decoder_cls 

5from evutils.types import TriggerArray 

6 

7class EventStreamer: 

8 """A low-level streamer that yields raw event blocks exactly as the parser outputs them. 

9 It performs ZERO slicing or logic; it purely drives the decoder. 

10 """ 

11 def __init__(self, file: "Path | str | bytes | io.BufferedIOBase", ext_trigger: bool = False, decoder_cls: "Type | None" = None, **kwargs) -> None: 

12 self.source = make_source(file) 

13 if decoder_cls is None: 

14 decoder_cls = resolve_decoder_cls(self.source) 

15 

16 # Set the flag as an attribute after construction (as EventReader does) 

17 # rather than forcing it through every decoder's __init__: several 

18 # decoders (DAT/NPZ/HDF5/CSV) don't accept it as a constructor kwarg. 

19 self.decoder = decoder_cls(self.source, **kwargs) 

20 self.decoder.read_external_triggers = ext_trigger 

21 self.ext_trigger = ext_trigger 

22 

23 def __iter__(self) -> "Iterator[EventArray]": 

24 self.decoder.init() 

25 while True: 

26 chunk = self.decoder.read_chunk() 

27 

28 if self.ext_trigger: 

29 if isinstance(chunk, tuple): 

30 ev_chunk, tr_chunk = chunk 

31 else: 

32 ev_chunk = chunk 

33 tr_chunk = TriggerArray.empty() 

34 else: 

35 ev_chunk = chunk[0] if isinstance(chunk, tuple) else chunk 

36 tr_chunk = None 

37 

38 if len(ev_chunk) == 0 and (tr_chunk is None or len(tr_chunk) == 0): 

39 break 

40 

41 if self.ext_trigger: 

42 yield ev_chunk, tr_chunk 

43 else: 

44 yield ev_chunk 

45