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

1"""Encoders module. 

2 

3Provides mapping and retrieval of event encoders based on file extensions. 

4Backends whose optional dependencies are missing (e.g. pandas for CSV, h5py 

5for HDF5) are not registered; asking for them raises an ``ImportError`` that 

6names the extra to install. 

7""" 

8 

9from pathlib import Path 

10from typing import Type 

11 

12from .common import EventEncoder 

13 

14#: Extension -> encoder class, for available backends only. 

15_WRITER_MAPPING: dict[str, Type[EventEncoder]] = {} 

16 

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

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

19 

20 

21from ._aedat import EventEncoder_Aedat 

22_WRITER_MAPPING[".aedat"] = EventEncoder_Aedat 

23 

24from ._bin import EventEncoder_Bin 

25_WRITER_MAPPING[".bin"] = EventEncoder_Bin 

26 

27try: 

28 from ._csv import EventEncoder_Csv 

29 _WRITER_MAPPING[".csv"] = EventEncoder_Csv 

30 _WRITER_MAPPING[".txt"] = EventEncoder_Csv 

31except ImportError: 

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

33 "writing CSV/TXT event files requires the evutils native library " 

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

35 ) 

36 

37from ._dat import EventEncoder_Dat 

38_WRITER_MAPPING[".dat"] = EventEncoder_Dat 

39 

40try: 

41 from ._hdf5 import EventEncoder_HDF5 

42 _WRITER_MAPPING[".h5"] = EventEncoder_HDF5 

43 _WRITER_MAPPING[".hdf5"] = EventEncoder_HDF5 

44except ImportError: 

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

46 "writing HDF5 event files requires h5py/hdf5plugin: install `evutils[hdf5]`" 

47 ) 

48 

49from ._npz import EventEncoder_Npz 

50_WRITER_MAPPING[".npz"] = EventEncoder_Npz 

51 

52from ._evt import EventEncoder_EVT 

53_WRITER_MAPPING[".raw"] = EventEncoder_EVT 

54_WRITER_MAPPING[".evt"] = EventEncoder_EVT 

55_WRITER_MAPPING[".evt3"] = EventEncoder_EVT 

56 

57from ._aer import EventEncoder_AER 

58_WRITER_MAPPING[".aer"] = EventEncoder_AER 

59 

60 

61def get_file_writer(file: Path) -> Type[EventEncoder]: 

62 """Get the appropriate writer for the given file. 

63 

64 Parameters 

65 ---------- 

66 file 

67 File to write 

68 

69 Returns 

70 ------- 

71 EventFileWriter 

72 Writer object for the file 

73 

74 """ 

75 ext = file.suffix.lower() 

76 if ext in _WRITER_MAPPING: 

77 return _WRITER_MAPPING[ext] 

78 if ext in _UNAVAILABLE: 

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

80 raise ValueError( 

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

82 f"{sorted(_WRITER_MAPPING.keys() | _UNAVAILABLE.keys())}" 

83 ) 

84 

85 

86__all__ = ["EventEncoder", "EventEncoder_Aedat", "EventEncoder_Bin", "EventEncoder_Csv", "EventEncoder_Dat", "EventEncoder_HDF5", "EventEncoder_Npz", "EventEncoder_EVT", "EventEncoder_AER", "get_file_writer"]