Coverage for larch/io/xas_data_source/spec.py: 100%

33 statements  

« prev     ^ index     » next       coverage.py v7.6.0, created at 2024-10-16 21:04 +0000

1import re 

2from typing import List, Optional, Iterator 

3from silx.io.specfile import SpecFile 

4from . import base 

5 

6 

7class SpecSingleXasDataSource(base.XasDataSource): 

8 """SPEC file. Each scan contains 1 XAS spectrum.""" 

9 

10 TYPE = "SPEC" 

11 

12 def __init__(self, *args, title_regex_pattern: Optional[str] = None, **kw) -> None: 

13 self.__specfile = None 

14 if title_regex_pattern: 

15 title_regex_pattern = re.compile(title_regex_pattern) 

16 self._title_regex_pattern = title_regex_pattern 

17 super().__init__(*args, **kw) 

18 

19 @property 

20 def _specfile(self): 

21 if self.__specfile is None: 

22 self.__specfile = SpecFile(self._filename) 

23 return self.__specfile 

24 

25 def get_source_info(self) -> str: 

26 return f"SPEC: {self._filename}" 

27 

28 def get_scan(self, scan_name: str) -> Optional[base.XasScan]: 

29 scan = self._specfile[scan_name] 

30 description = "\n".join(scan.header) 

31 return base.XasScan( 

32 name=scan_name, 

33 description=description, 

34 start_time="TODO", 

35 info=self.get_source_info(), 

36 labels=scan.labels, 

37 data=scan.data, 

38 ) 

39 

40 def get_scan_names(self) -> List[str]: 

41 return list(self._iter_scan_names()) 

42 

43 def _iter_scan_names(self) -> Iterator[str]: 

44 for scan in self._specfile: 

45 if self._title_regex_pattern is not None: 

46 title = scan.scan_header_dict["S"] 

47 title = "".join(title.split(" ")[1:]) 

48 if not self._title_regex_pattern.match(title): 

49 continue 

50 yield f"{scan.number}.{scan.order}"