Coverage for src/susi/db/metadata.py: 97%

1from __future__ import annotations 

2 

3import os 

4import pickle 

5import re 

6from typing import Union 

7 

8from ..base import Logging 

9from ..base.header_keys import * 

10from ..io import Fits 

11 

12log = Logging.get_logger() 

13 

14 

15class Metadata: 

16 """ 

17 Metadata database based on pickel files that contain a python dict. 

18 """ 

19 

20 DB_KEYS = [ 

21 DATE_OBS, 

22 TIMESTAMP_US, 

23 CAMERA_NAME, 

24 CAMERA_ID, 

25 INTEGRATION_TIME, 

26 HK_PMU_CUR, 

27 HK_PMU_ANG, 

28 HK_PMU_ANG_ERR, 

29 MOD_STATE, 

30 IMG_RMS, 

31 IMG_SNR, 

32 IMG_MEAN, 

33 IMG_CONTRAST, 

34 ADC15_CF_TEMP, 

35 ROI_X0, 

36 ROI_X1, 

37 ROI_Y0, 

38 ROI_Y1, 

39 TEM7, 

40 SHEAR_CORR, 

41 ROTATION_ANG, 

42 XCEN, 

43 YCEN, 

44 CW_LOOP, 

45 FLAT_MOD, 

46 PS_STATE, 

47 SLIT_FLAT_OFFSET, 

48 SLIT_FLAT_SLOPE, 

49 ] 

50 

51 @staticmethod 

52 def insert(fits: Fits, overwrite=False) -> None: 

53 """ 

54 Insert the metadata from the given fits file to the appropriate database. 

55 

56 The database is selected by the full path. 

57 For each hourly folder there should be a database file with a matching name. 

58 

59 All selected metadata is added to the database and is committed (saved) afterwards automatically 

60 """ 

61 db = Metadata.db_path(fits) 

62 if db is not None: 

63 Metadata(db).add(fits, overwrite=overwrite).commit() 

64 

65 @staticmethod 

66 def insert_batch(fitses: list, overwrite=False) -> None: 

67 """ 

68 Insert the metadata from the given list of fits file to the appropriate database. 

69 Note The list is assumed to be sorted from oldest to newest! 

70 

71 The database is selected by the full path of the first fits file. 

72 For each hourly folder there should be a database file with a matching name. 

73 

74 All selected metadata is added to the database and is committed (saved) afterwards automatically 

75 """ 

76 db0 = Metadata.db_path(fitses[0]) 

77 db1 = Metadata.db_path(fitses[1]) 

78 if db0 is None or db0 != db1: 

79 raise ValueError('Cannot save the metadata of the batch to a single DB file...') 

80 

81 db = Metadata(db0) 

82 for file in fitses: 

83 db.add(file, overwrite=overwrite) 

84 db.commit() 

85 

86 @staticmethod 

87 def db_path(fits: Union[Fits, str]) -> Union[str, None]: 

88 """ 

89 Get the matching metadata file for the given fits 

90 

91 The database is selected by the full path. 

92 For each hourly folder there should be a database file with a matching name. 

93 """ 

94 path = fits.path if isinstance(fits, Fits) else fits 

95 m = re.match(r".*(\d{4}_\d{2}_\d{2}_\d{2}).*(cam\d).*", path) 

96 if m: 

97 filename = f'{m.group(1)}-{m.group(2)}.pickle' 

98 return os.path.join(path.split(m.group(1))[0], filename) 

99 

100 m = re.match(r".*(\d{4}_\d{2}_\d{2}_\d{2}).*", path) 

101 if m: 

102 filename = f'{m.group(1)}.pickle' 

103 return os.path.join(path.split(m.group(1))[0], filename) 

104 

105 log.error('Could not determine database by fits file name %s', path) 

106 return None 

107 

108 def __init__(self, path: str): 

109 self.changed = False 

110 self.path = path 

111 self.data = {} 

112 self.reload() 

113 

114 def reload(self) -> None: 

115 if not os.path.isfile(self.path): 

116 return 

117 

118 with open(self.path, 'rb') as f: 

119 self.data = pickle.load(f) 

120 

121 def __contains__(self, file: str) -> bool: 

122 """Check if a given entry is in the database""" 

123 return os.path.basename(file) in self.data 

124 

125 def __getitem__(self, file: str) -> Union[dict, None]: 

126 key = os.path.basename(file) 

127 return self.data[key] if key in self else None 

128 

129 def add(self, fits: Fits, overwrite=False) -> Metadata: 

130 """Add an entry to the metadata database""" 

131 key = os.path.basename(fits.path) 

132 if overwrite or key not in self: 

133 self.changed = True 

134 self.data[key] = Metadata._generate_entry(fits) 

135 return self 

136 

137 def delete(self, file: Union[Fits, str]): 

138 """Remove an entry from the metadata database""" 

139 path = file.path if isinstance(file, Fits) else file 

140 self.data.pop(os.path.basename(path), None) 

141 

142 def commit(self) -> None: 

143 """Write all changes (if any) back to disk""" 

144 if not self.changed: 

145 return 

146 

147 os.makedirs(os.path.dirname(self.path), exist_ok=True) 

148 with open(self.path, 'wb', 0o775) as f: 

149 pickle.dump(self.data, f) 

150 self.changed = False 

151 

152 @staticmethod 

153 def _generate_entry(fits: Fits): 

154 entry = {} 

155 for key in Metadata.DB_KEYS: 

156 try: 

157 entry[key] = fits.header[key] 

158 except KeyError: 

159 entry[key] = None 

160 return entry