1from pathlib import Path 

2import json 

3from collections.abc import MutableMapping, Mapping 

4from contextlib import contextmanager 

5from ase.io.jsonio import read_json, write_json, encode 

6from ase.utils import opencew 

7 

8 

9def missing(key): 

10 raise KeyError(key) 

11 

12 

13class Locked(Exception): 

14 pass 

15 

16 

17class CacheLock: 

18 def __init__(self, fd, key): 

19 self.fd = fd 

20 self.key = key 

21 

22 def save(self, value): 

23 json_utf8 = encode(value).encode('utf-8') 

24 try: 

25 self.fd.write(json_utf8) 

26 except Exception as ex: 

27 raise RuntimeError(f'Failed to save {value} to cache') from ex 

28 finally: 

29 self.fd.close() 

30 

31 

32class MultiFileJSONCache(MutableMapping): 

33 writable = True 

34 

35 def __init__(self, directory): 

36 self.directory = Path(directory) 

37 

38 def _filename(self, key): 

39 return self.directory / f'cache.{key}.json' 

40 

41 def _glob(self): 

42 return self.directory.glob('cache.*.json') 

43 

44 def __iter__(self): 

45 for path in self._glob(): 

46 cache, key = path.stem.split('.', 1) 

47 if cache != 'cache': 

48 continue 

49 yield key 

50 

51 def __len__(self): 

52 # Very inefficient this, but not a big usecase. 

53 return len(list(self._glob())) 

54 

55 @contextmanager 

56 def lock(self, key): 

57 self.directory.mkdir(exist_ok=True, parents=True) 

58 path = self._filename(key) 

59 fd = opencew(path) 

60 try: 

61 if fd is None: 

62 yield None 

63 else: 

64 yield CacheLock(fd, key) 

65 finally: 

66 if fd is not None: 

67 fd.close() 

68 

69 def __setitem__(self, key, value): 

70 with self.lock(key) as handle: 

71 if handle is None: 

72 raise Locked(key) 

73 handle.save(value) 

74 

75 def __getitem__(self, key): 

76 path = self._filename(key) 

77 try: 

78 return read_json(path, always_array=False) 

79 except FileNotFoundError: 

80 missing(key) 

81 except json.decoder.JSONDecodeError: 

82 # May be partially written, which typically means empty 

83 # because the file was locked with exclusive-write-open. 

84 # 

85 # Since we decide what keys we have based on which files exist, 

86 # we are obligated to return a value for this case too. 

87 # So we return None. 

88 return None 

89 

90 def __delitem__(self, key): 

91 try: 

92 self._filename(key).unlink() 

93 except FileNotFoundError: 

94 missing(key) 

95 

96 def combine(self): 

97 cache = CombinedJSONCache.dump_cache(self.directory, dict(self)) 

98 assert set(cache) == set(self) 

99 self.clear() 

100 assert len(self) == 0 

101 return cache 

102 

103 def split(self): 

104 return self 

105 

106 def filecount(self): 

107 return len(self) 

108 

109 def strip_empties(self): 

110 empties = [key for key, value in self.items() if value is None] 

111 for key in empties: 

112 del self[key] 

113 return len(empties) 

114 

115 

116class CombinedJSONCache(Mapping): 

117 writable = False 

118 

119 def __init__(self, directory, dct): 

120 self.directory = Path(directory) 

121 self._dct = dict(dct) 

122 

123 def filecount(self): 

124 return int(self._filename.is_file()) 

125 

126 @property 

127 def _filename(self): 

128 return self.directory / 'combined.json' 

129 

130 def _dump_json(self): 

131 target = self._filename 

132 if target.exists(): 

133 raise RuntimeError(f'Already exists: {target}') 

134 self.directory.mkdir(exist_ok=True, parents=True) 

135 write_json(target, self._dct) 

136 

137 def __len__(self): 

138 return len(self._dct) 

139 

140 def __iter__(self): 

141 return iter(self._dct) 

142 

143 def __getitem__(self, index): 

144 return self._dct[index] 

145 

146 @classmethod 

147 def dump_cache(cls, path, dct): 

148 cache = cls(path, dct) 

149 cache._dump_json() 

150 return cache 

151 

152 @classmethod 

153 def load(cls, path): 

154 # XXX Very hacky this one 

155 cache = cls(path, {}) 

156 dct = read_json(cache._filename, always_array=False) 

157 cache._dct.update(dct) 

158 return cache 

159 

160 def clear(self): 

161 self._filename.unlink() 

162 self._dct.clear() 

163 

164 def combine(self): 

165 return self 

166 

167 def split(self): 

168 cache = MultiFileJSONCache(self.directory) 

169 assert len(cache) == 0 

170 cache.update(self) 

171 assert set(cache) == set(self) 

172 self.clear() 

173 return cache 

174 

175 

176def get_json_cache(directory): 

177 try: 

178 return CombinedJSONCache.load(directory) 

179 except FileNotFoundError: 

180 return MultiFileJSONCache(directory)