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

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 

13from ._compression import strip_compression_suffix 

14 

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

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

17 

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

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

20 

21from ._aedat import EventEncoder_Aedat 

22_WRITER_MAPPING[".aedat"] = EventEncoder_Aedat 

23_WRITER_MAPPING[".aedat4"] = EventEncoder_Aedat 

24 

25from ._bin import EventEncoder_Bin 

26_WRITER_MAPPING[".bin"] = EventEncoder_Bin 

27 

28try: 

29 from ._csv import EventEncoder_Csv 

30 _WRITER_MAPPING[".csv"] = EventEncoder_Csv 

31 _WRITER_MAPPING[".txt"] = EventEncoder_Csv 

32except ImportError: 

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

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

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

36 ) 

37 

38from ._dat import EventEncoder_Dat 

39_WRITER_MAPPING[".dat"] = EventEncoder_Dat 

40 

41try: 

42 from ._hdf5 import EventEncoder_HDF5 

43 _WRITER_MAPPING[".h5"] = EventEncoder_HDF5 

44 _WRITER_MAPPING[".hdf5"] = EventEncoder_HDF5 

45except ImportError: 

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

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

48 ) 

49 

50from ._npz import EventEncoder_Npz 

51_WRITER_MAPPING[".npz"] = EventEncoder_Npz 

52 

53from ._evt import EventEncoder_EVT 

54 

55def _evt_encoder_for(fmt: str) -> Type[EventEncoder]: 

56 """EventEncoder_EVT preconfigured for ``fmt``; an explicit ``format=`` 

57 kwarg from the caller still wins.""" 

58 class _Encoder(EventEncoder_EVT): 

59 def __init__(self, *args, format: str = fmt, **kwargs): # noqa: A002 

60 super().__init__(*args, format=format, **kwargs) 

61 _Encoder.__name__ = f"EventEncoder_EVT_{fmt}" 

62 _Encoder.__qualname__ = _Encoder.__name__ 

63 return _Encoder 

64 

65_WRITER_MAPPING[".raw"] = EventEncoder_EVT 

66_WRITER_MAPPING[".evt"] = EventEncoder_EVT 

67_WRITER_MAPPING[".evt3"] = EventEncoder_EVT 

68_WRITER_MAPPING[".evt2"] = _evt_encoder_for("evt2") 

69_WRITER_MAPPING[".evt21"] = _evt_encoder_for("evt21") 

70_WRITER_MAPPING[".evt4"] = _evt_encoder_for("evt4") 

71 

72from ._aer import EventEncoder_AER 

73_WRITER_MAPPING[".aer"] = EventEncoder_AER 

74 

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

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

77 

78 Parameters 

79 ---------- 

80 file 

81 File to write 

82 

83 Returns 

84 ------- 

85 EventFileWriter 

86 Writer object for the file 

87 

88 """ 

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

90 # *inner* extension selects the encoder. 

91 ext = Path(strip_compression_suffix(str(file))).suffix.lower() 

92 if ext in _WRITER_MAPPING: 

93 return _WRITER_MAPPING[ext] 

94 if ext in _UNAVAILABLE: 

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

96 raise ValueError( 

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

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

99 ) 

100 

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