Coverage for src/evutils/io/decoders.py: 82%

1"""Decoders module. 

2 

3Provides mapping and resolution of event decoders based on file extensions or 

4magic bytes. Backends whose optional dependencies are missing (e.g. pandas for 

5CSV, h5py for HDF5) are not registered; asking for them raises an 

6``ImportError`` that names the extra to install. 

7""" 

8 

9from pathlib import Path 

10from typing import Type, cast 

11 

12from .common import EventDecoder 

13from ._compression import strip_compression_suffix 

14 

15#: Extension -> decoder class, for available backends only. 

16_READER_MAPPING: dict[str, Type[EventDecoder]] = {} 

17 

18#: Extension -> reason it is unavailable (missing optional dependency). 

19_UNAVAILABLE: dict[str, str] = {} 

20 

21from ._aedat import EventDecoder_Aedat 

22_READER_MAPPING[".aedat"] = EventDecoder_Aedat 

23_READER_MAPPING[".aedat4"] = EventDecoder_Aedat 

24 

25from ._bin import EventDecoder_Bin 

26_READER_MAPPING[".bin"] = EventDecoder_Bin 

27 

28try: 

29 from ._csv import EventDecoder_Csv 

30 _READER_MAPPING[".csv"] = EventDecoder_Csv 

31 _READER_MAPPING[".txt"] = EventDecoder_Csv 

32except ImportError: 

33 _UNAVAILABLE[".csv"] = _UNAVAILABLE[".txt"] = ( 

34 "reading CSV/TXT event files requires the evutils native library " 

35 "(build it with `uv pip install -e .`)" 

36 ) 

37 

38from ._dat import EventDecoder_Dat 

39_READER_MAPPING[".dat"] = EventDecoder_Dat 

40 

41try: 

42 from ._hdf5 import EventDecoder_HDF5 

43 _READER_MAPPING[".hdf5"] = EventDecoder_HDF5 

44 _READER_MAPPING[".h5"] = EventDecoder_HDF5 

45except ImportError: 

46 _UNAVAILABLE[".hdf5"] = _UNAVAILABLE[".h5"] = ( 

47 "reading HDF5 event files requires h5py/hdf5plugin: install `evutils[hdf5]`" 

48 ) 

49 

50from ._npz import EventDecoder_Npz 

51_READER_MAPPING[".npz"] = EventDecoder_Npz 

52 

53from ._evt import EventDecoder_EVT 

54_READER_MAPPING[".raw"] = EventDecoder_EVT 

55_READER_MAPPING[".evt"] = EventDecoder_EVT 

56_READER_MAPPING[".evt3"] = EventDecoder_EVT 

57_READER_MAPPING[".evt2"] = EventDecoder_EVT 

58_READER_MAPPING[".evt21"] = EventDecoder_EVT 

59_READER_MAPPING[".evt4"] = EventDecoder_EVT 

60 

61from ._aer import EventDecoder_AER 

62_READER_MAPPING[".aer"] = EventDecoder_AER 

63 

64def _lookup(ext: str) -> Type[EventDecoder]: 

65 """Resolve an extension to a decoder class, or raise a helpful error.""" 

66 if ext in _READER_MAPPING: 

67 return _READER_MAPPING[ext] 

68 if ext in _UNAVAILABLE: 

69 raise ImportError(f"File extension {ext} is supported, but {_UNAVAILABLE[ext]}") 

70 raise ValueError( 

71 f"File extension {ext} not supported, available extensions: " 

72 f"{sorted(_READER_MAPPING.keys() | _UNAVAILABLE.keys())}" 

73 ) 

74 

75def get_reader_from_filename(file: Path) -> Type[EventDecoder]: 

76 """Get the appropriate reader for the given file. 

77 

78 Parameters 

79 ---------- 

80 file 

81 File to read 

82 

83 Returns 

84 ------- 

85 EventDecoder 

86 Reader object for the file 

87 

88 """ 

89 return _lookup(file.suffix.lower()) 

90 

91# Content sniffers: (predicate over the first bytes -> decoder class). Tried in 

92# order when the filename extension is unknown or absent (streams, USB). 

93def _header_lines(head: bytes) -> list[str]: 

94 """The ``"% ..."`` ASCII header lines at the start of ``head``.""" 

95 text = head.decode("ascii", "ignore") 

96 return [ln for ln in text.split("\n") if ln.startswith("% ")] 

97 

98def _sniff_dat(head: bytes) -> bool: 

99 """Prophesee DAT: header carries ``% Version`` / ``% Data file containing``. 

100 

101 DAT and RAW/EVT both open with a ``%`` header, so they can only be told 

102 apart by the keywords inside it. 

103 """ 

104 for ln in _header_lines(head): 

105 low = ln.lower() 

106 if "data file containing" in low or low.startswith("% version"): 

107 return True 

108 return False 

109 

110def _sniff_evt(head: bytes) -> bool: 

111 """Prophesee RAW/EVT: header carries ``% evt`` / ``% format EVT`` / ``% geometry``.""" 

112 for ln in _header_lines(head): 

113 low = ln.lower() 

114 if (low.startswith("% evt") 

115 or low.startswith("% format evt") 

116 or low.startswith("% geometry")): 

117 return True 

118 return False 

119 

120def _sniff_prophesee(head: bytes) -> bool: 

121 """Fallback: any other ``"% "``-headed stream with prophesee-like metadata.""" 

122 if not head.startswith(b"% "): 

123 return False 

124 low = head.lower() 

125 return b"system_id" in low or b"firmware_version" in low or b"plugin name" in low 

126 

127def _sniff_aedat(head: bytes) -> bool: 

128 """Check if the first bytes match an AEDAT version line. 

129 

130 Parameters 

131 ---------- 

132 head : bytes 

133 The first bytes of the file/stream. 

134 

135 Returns 

136 ------- 

137 bool 

138 True if it matches the AEDAT format, False otherwise. 

139 

140 """ 

141 return head.startswith(b"#!AER-DAT") 

142 

143_SNIFFERS = [ 

144 (_sniff_dat, "EventDecoder_Dat"), 

145 (_sniff_evt, "EventDecoder_EVT"), 

146 (_sniff_aedat, "EventDecoder_Aedat"), 

147 (_sniff_prophesee, "EventDecoder_EVT"), # generic "% "-headed fallback 

148] 

149 

150def resolve_decoder_cls(source: "io.BufferedIOBase | str | bytes") -> Type[EventDecoder]: 

151 """Determine the decoder class for a :class:`ByteSource`. 

152 

153 Tries the filename extension first (cheap, usually right), then falls back 

154 to sniffing the leading bytes -- which works for streams and USB devices 

155 that have no filename. 

156 

157 Parameters 

158 ---------- 

159 source 

160 A ByteSource (see :mod:`evutils.io._source`). 

161 

162 Returns 

163 ------- 

164 Type[EventDecoder] 

165 The decoder class to instantiate with the source. 

166 

167 """ 

168 name = getattr(source, "name", None) 

169 if name: 

170 # A compression suffix (foo.raw.zst) is transparent to format choice: 

171 # strip it so the *inner* extension selects the decoder. 

172 name = strip_compression_suffix(name) 

173 ext = Path(name).suffix.lower() 

174 if ext in _READER_MAPPING or ext in _UNAVAILABLE: 

175 return _lookup(ext) 

176 

177 try: 

178 head = source.peek(512) 

179 except Exception: 

180 head = b"" 

181 

182 for matches, cls_name in _SNIFFERS: 

183 if matches(head): 

184 return cast(Type[EventDecoder], globals()[cls_name]) 

185 

186 raise ValueError( 

187 "Could not determine the event format: unknown/absent extension " 

188 f"({name!r}) and no known magic bytes. Pass an explicit decoder." 

189 ) 

190 

191__all__ = ["EventDecoder", 'EventDecoder_Aedat', 'EventDecoder_Bin', 'EventDecoder_Csv', 'EventDecoder_Dat', 'EventDecoder_HDF5', 'EventDecoder_Npz', 'EventDecoder_EVT', 'EventDecoder_AER', 'get_reader_from_filename', 'resolve_decoder_cls']