1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
|
from pathlib import Path
import json
from collections.abc import MutableMapping, Mapping
from contextlib import contextmanager
from ase.io.jsonio import read_json, write_json, encode
from ase.utils import opencew
def missing(key):
raise KeyError(key)
class Locked(Exception):
pass
class CacheLock:
def __init__(self, fd, key):
self.fd = fd
self.key = key
def save(self, value):
json_utf8 = encode(value).encode('utf-8')
try:
self.fd.write(json_utf8)
except Exception as ex:
raise RuntimeError(f'Failed to save {value} to cache') from ex
finally:
self.fd.close()
class MultiFileJSONCache(MutableMapping):
writable = True
def __init__(self, directory):
self.directory = Path(directory)
def _filename(self, key):
return self.directory / f'cache.{key}.json'
def _glob(self):
return self.directory.glob('cache.*.json')
def __iter__(self):
for path in self._glob():
cache, key = path.stem.split('.', 1)
if cache != 'cache':
continue
yield key
def __len__(self):
# Very inefficient this, but not a big usecase.
return len(list(self._glob()))
@contextmanager
def lock(self, key):
self.directory.mkdir(exist_ok=True, parents=True)
path = self._filename(key)
fd = opencew(path)
try:
if fd is None:
yield None
else:
yield CacheLock(fd, key)
finally:
if fd is not None:
fd.close()
def __setitem__(self, key, value):
with self.lock(key) as handle:
if handle is None:
raise Locked(key)
handle.save(value)
def __getitem__(self, key):
path = self._filename(key)
try:
return read_json(path, always_array=False)
except FileNotFoundError:
missing(key)
except json.decoder.JSONDecodeError:
# May be partially written, which typically means empty
# because the file was locked with exclusive-write-open.
#
# Since we decide what keys we have based on which files exist,
# we are obligated to return a value for this case too.
# So we return None.
return None
def __delitem__(self, key):
try:
self._filename(key).unlink()
except FileNotFoundError:
missing(key)
def combine(self):
cache = CombinedJSONCache.dump_cache(self.directory, dict(self))
assert set(cache) == set(self)
self.clear()
assert len(self) == 0
return cache
def split(self):
return self
def filecount(self):
return len(self)
def strip_empties(self):
empties = [key for key, value in self.items() if value is None]
for key in empties:
del self[key]
return len(empties)
class CombinedJSONCache(Mapping):
writable = False
def __init__(self, directory, dct):
self.directory = Path(directory)
self._dct = dict(dct)
def filecount(self):
return int(self._filename.is_file())
@property
def _filename(self):
return self.directory / 'combined.json'
def _dump_json(self):
target = self._filename
if target.exists():
raise RuntimeError(f'Already exists: {target}')
self.directory.mkdir(exist_ok=True, parents=True)
write_json(target, self._dct)
def __len__(self):
return len(self._dct)
def __iter__(self):
return iter(self._dct)
def __getitem__(self, index):
return self._dct[index]
@classmethod
def dump_cache(cls, path, dct):
cache = cls(path, dct)
cache._dump_json()
return cache
@classmethod
def load(cls, path):
# XXX Very hacky this one
cache = cls(path, {})
dct = read_json(cache._filename, always_array=False)
cache._dct.update(dct)
return cache
def clear(self):
self._filename.unlink()
self._dct.clear()
def combine(self):
return self
def split(self):
cache = MultiFileJSONCache(self.directory)
assert len(cache) == 0
cache.update(self)
assert set(cache) == set(self)
self.clear()
return cache
def get_json_cache(directory):
try:
return CombinedJSONCache.load(directory)
except FileNotFoundError:
return MultiFileJSONCache(directory)
|
