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

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 db = Metadata(db0) 

81 for file in fitses: 

82 db.add(file, overwrite=overwrite) 

83 db.commit() 

84 

85 @staticmethod 

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

87 """ 

88 Get the matching metadata file for the given fits 

89 

90 The database is selected by the full path. 

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

92 """ 

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

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

95 if m: 

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

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

98 

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

100 if m: 

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

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

103 

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

105 return None 

106 

107 def __init__(self, path: str): 

108 self.changed = False 

109 self.path = path 

110 self.data = {} 

111 self.reload() 

112 

113 def reload(self) -> None: 

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

115 return 

116 

117 try: 

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

119 self.data = pickle.load(f) 

120 except (pickle.UnpicklingError, EOFError, FileNotFoundError) as e: 

121 log.error("Error loading metadata file %s: %s", self.path, e) 

122 log.error("Deleting corrupted file and creating a new one.") 

123 os.remove(self.path) 

124 self.data = {} 

125 return 

126 

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

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

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

130 

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

132 key = os.path.basename(file) 

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

134 

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

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

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

138 if overwrite or key not in self: 

139 self.changed = True 

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

141 return self 

142 

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

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

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

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

147 

148 def commit(self) -> None: 

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

150 if not self.changed: 

151 return 

152 

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

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

155 pickle.dump(self.data, f) 

156 self.changed = False 

157 

158 @staticmethod 

159 def _generate_entry(fits: Fits): 

160 entry = {} 

161 for key in Metadata.DB_KEYS: 

162 try: 

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

164 except KeyError: 

165 entry[key] = None 

166 return entry