Coverage for src/finbot/market/cache.py: 96%

1import hashlib 

2import json 

3import logging 

4import shutil 

5import tempfile 

6from collections.abc import Mapping, Sequence 

7from datetime import datetime 

8from pathlib import Path 

9 

10from finbot.market.dataset import DatasetMetadata 

11from finbot.market.models import Bar 

12 

13logger = logging.getLogger(__name__) 

14 

15CACHE_FORMAT_VERSION = 1 

16DEFAULT_CACHE_DIR = Path(".finbot/cache") 

17 

18 

19class DatasetCacheError(RuntimeError): 

20 """Base error for historical dataset cache failures.""" 

21 

22 

23class DatasetNotFoundError(DatasetCacheError): 

24 """Raised when a requested dataset is not present in the cache.""" 

25 

26 

27class DatasetIntegrityError(DatasetCacheError): 

28 """Raised when cached dataset files are incomplete or inconsistent.""" 

29 

30 

31class HistoricalDataCache: 

32 def __init__(self, root: str | Path = DEFAULT_CACHE_DIR): 

33 self.root = Path(root) 

34 

35 def contains(self, dataset_id: str) -> bool: 

36 return self._dataset_dir(dataset_id).is_dir() 

37 

38 def save(self, metadata: DatasetMetadata, bars: Mapping[str, Sequence[Bar]]) -> str: 

39 logger.debug("Preparing dataset %s for cache", metadata.dataset_id) 

40 

41 normalized_bars = self._normalize_bars(metadata, bars) 

42 bars_bytes = self._serialize_bars(normalized_bars) 

43 bars_sha256 = hashlib.sha256(bars_bytes).hexdigest() 

44 

45 dataset_id = metadata.dataset_id 

46 

47 bar_count = sum(len(items) for items in normalized_bars.values()) 

48 logger.debug( 

49 "Dataset %s contains %d bars; SHA-256=%s", 

50 dataset_id, 

51 bar_count, 

52 bars_sha256, 

53 ) 

54 

55 dataset_dir = self._dataset_dir(dataset_id) 

56 

57 if dataset_dir.exists(): 

58 cached_metadata, cached_bars = self.load(dataset_id) 

59 

60 logger.debug( 

61 "Dataset %s already exists in cache; validating existing files", 

62 dataset_id, 

63 ) 

64 

65 if cached_metadata != metadata: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true

66 raise DatasetIntegrityError( 

67 f"Cached dataset {dataset_id!r} has metadata " 

68 "that does not match its dataset ID.", 

69 ) 

70 

71 if self._serialize_bars(cached_bars) != bars_bytes: 

72 raise DatasetIntegrityError( 

73 f"Dataset {dataset_id!r} already exists with different bars.", 

74 ) 

75 

76 logger.debug("Existing cached dataset %s matches requested dataset", dataset_id) 

77 return dataset_id 

78 

79 logger.debug("Writing dataset %s to cache directory %s", dataset_id, dataset_dir) 

80 

81 self.root.mkdir(parents=True, exist_ok=True) 

82 

83 temp_dir = Path(tempfile.mkdtemp(prefix=f".{dataset_id}.", dir=self.root)) 

84 

85 try: 

86 metadata_bytes = self._json_bytes(metadata.to_dict()) 

87 

88 manifest = { 

89 "cache_format_version": CACHE_FORMAT_VERSION, 

90 "dataset_id": dataset_id, 

91 "bar_count": sum(len(items) for items in normalized_bars.values()), 

92 "bars_sha256": bars_sha256, 

93 } 

94 

95 (temp_dir / "metadata.json").write_bytes(metadata_bytes) 

96 (temp_dir / "bars.jsonl").write_bytes(bars_bytes) 

97 (temp_dir / "manifest.json").write_bytes(self._json_bytes(manifest)) 

98 

99 temp_dir.rename(dataset_dir) 

100 

101 logger.debug("Dataset %s cache write complete", dataset_id) 

102 

103 except Exception: 

104 shutil.rmtree(temp_dir, ignore_errors=True) 

105 raise 

106 

107 return dataset_id 

108 

109 def load(self, dataset_id: str) -> tuple[DatasetMetadata, dict[str, list[Bar]]]: 

110 dataset_dir = self._dataset_dir(dataset_id) 

111 

112 if not dataset_dir.is_dir(): 

113 raise DatasetNotFoundError( 

114 f"Dataset {dataset_id!r} was not found in cache {self.root}.", 

115 ) 

116 

117 logger.debug("Loading cached dataset %s from %s", dataset_id, dataset_dir) 

118 

119 metadata_path = dataset_dir / "metadata.json" 

120 bars_path = dataset_dir / "bars.jsonl" 

121 manifest_path = dataset_dir / "manifest.json" 

122 

123 for path in (metadata_path, bars_path, manifest_path): 

124 if not path.is_file(): 

125 raise DatasetIntegrityError( 

126 f"Cached dataset {dataset_id!r} is incomplete: missing {path.name}.", 

127 ) 

128 

129 metadata_data = self._read_json_object(metadata_path, dataset_id) 

130 

131 try: 

132 metadata = DatasetMetadata.from_dict(metadata_data) 

133 except (TypeError, ValueError) as exc: 

134 raise DatasetIntegrityError( 

135 f"Cached dataset {dataset_id!r} has invalid metadata: {exc}", 

136 ) from exc 

137 

138 if metadata.dataset_id != dataset_id: 

139 raise DatasetIntegrityError( 

140 f"Cached dataset {dataset_id!r} has metadata for dataset {metadata.dataset_id!r}.", 

141 ) 

142 

143 manifest = self._read_json_object(manifest_path, dataset_id) 

144 self._validate_manifest(manifest, dataset_id) 

145 

146 bars_bytes = bars_path.read_bytes() 

147 actual_hash = hashlib.sha256(bars_bytes).hexdigest() 

148 expected_hash = manifest["bars_sha256"] 

149 

150 logger.debug("Validating bars checksum for dataset %s", dataset_id) 

151 

152 if actual_hash != expected_hash: 

153 raise DatasetIntegrityError( 

154 f"Cached dataset {dataset_id!r} failed checksum validation for bars.jsonl.", 

155 ) 

156 

157 bars = self._deserialize_bars(bars_bytes, dataset_id) 

158 normalized_bars = self._normalize_bars(metadata, bars) 

159 actual_count = sum(len(items) for items in normalized_bars.values()) 

160 

161 if actual_count != manifest["bar_count"]: 

162 raise DatasetIntegrityError( 

163 f"Cached dataset {dataset_id!r} has {actual_count} bars, " 

164 f"but its manifest declares {manifest['bar_count']}.", 

165 ) 

166 

167 logger.debug("Cached dataset %s validated successfully: %d bars", dataset_id, actual_count) 

168 

169 return metadata, normalized_bars 

170 

171 def get_bars(self, dataset_id: str) -> dict[str, list[Bar]]: 

172 _, bars = self.load(dataset_id) 

173 return bars 

174 

175 def get_metadata(self, dataset_id: str) -> DatasetMetadata: 

176 metadata, _ = self.load(dataset_id) 

177 return metadata 

178 

179 def _dataset_dir(self, dataset_id: str) -> Path: 

180 if not dataset_id or any(char not in "0123456789abcdef" for char in dataset_id): 

181 raise ValueError("Dataset ID must be a non-empty lowercase hexadecimal string.") 

182 

183 return self.root / dataset_id 

184 

185 @staticmethod 

186 def _json_bytes(data: object) -> bytes: 

187 payload = json.dumps(data, sort_keys=True, separators=(",", ":")) 

188 return (payload + "\n").encode("utf-8") 

189 

190 @classmethod 

191 def _serialize_bars(cls, bars: Mapping[str, Sequence[Bar]]) -> bytes: 

192 records: list[bytes] = [] 

193 

194 for symbol in sorted(bars): 

195 for bar in bars[symbol]: 

196 records.append( 

197 cls._json_bytes( 

198 { 

199 "symbol": bar.symbol, 

200 "timestamp": bar.timestamp.isoformat(), 

201 "open": float(bar.open), 

202 "high": float(bar.high), 

203 "low": float(bar.low), 

204 "close": float(bar.close), 

205 "volume": float(bar.volume), 

206 }, 

207 ), 

208 ) 

209 

210 return b"".join(records) 

211 

212 @staticmethod 

213 def _deserialize_bars(payload: bytes, dataset_id: str) -> dict[str, list[Bar]]: 

214 result: dict[str, list[Bar]] = {} 

215 

216 for line_number, raw_line in enumerate(payload.splitlines(), start=1): 

217 if not raw_line.strip(): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true

218 continue 

219 

220 try: 

221 data = json.loads(raw_line) 

222 

223 if not isinstance(data, dict): 

224 raise ValueError("bar record must be a JSON object") 

225 

226 bar = Bar( 

227 symbol=data["symbol"], 

228 timestamp=datetime.fromisoformat(data["timestamp"]), 

229 open=float(data["open"]), 

230 high=float(data["high"]), 

231 low=float(data["low"]), 

232 close=float(data["close"]), 

233 volume=float(data["volume"]), 

234 ) 

235 

236 except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: 

237 raise DatasetIntegrityError( 

238 f"Cached dataset {dataset_id!r} contains an " 

239 f"invalid bar at line {line_number}: {exc}", 

240 ) from exc 

241 

242 result.setdefault(bar.symbol, []).append(bar) 

243 

244 return result 

245 

246 @staticmethod 

247 def _normalize_bars( 

248 metadata: DatasetMetadata, 

249 bars: Mapping[str, Sequence[Bar]], 

250 ) -> dict[str, list[Bar]]: 

251 expected_symbols = set(metadata.symbols) 

252 actual_symbols = set(bars) 

253 

254 if actual_symbols != expected_symbols: 

255 missing = sorted(expected_symbols - actual_symbols) 

256 unexpected = sorted(actual_symbols - expected_symbols) 

257 

258 details: list[str] = [] 

259 

260 if missing: 260 ↛ 263line 260 didn't jump to line 263 because the condition on line 260 was always true

261 details.append(f"missing symbols: {", ".join(missing)}") 

262 

263 if unexpected: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true

264 details.append(f"unexpected symbols: {", ".join(unexpected)}") 

265 

266 raise DatasetIntegrityError( 

267 f"Bars do not match dataset metadata ({"; ".join(details)}).", 

268 ) 

269 

270 result: dict[str, list[Bar]] = {} 

271 

272 for symbol in metadata.symbols: 

273 items = list(bars[symbol]) 

274 

275 for bar in items: 

276 if bar.symbol != symbol: 

277 raise DatasetIntegrityError( 

278 f"Bar stored under {symbol!r} has symbol {bar.symbol!r}.", 

279 ) 

280 

281 if bar.timestamp.tzinfo is None: 

282 raise DatasetIntegrityError( 

283 f"Bar for {symbol!r} has a timezone-naive timestamp.", 

284 ) 

285 

286 result[symbol] = sorted(items, key=lambda bar: bar.timestamp) 

287 

288 return result 

289 

290 @staticmethod 

291 def _read_json_object(path: Path, dataset_id: str) -> dict[str, object]: 

292 try: 

293 data = json.loads(path.read_text(encoding="utf-8")) 

294 except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: 

295 raise DatasetIntegrityError( 

296 f"Cached dataset {dataset_id!r} has invalid {path.name}: {exc}", 

297 ) from exc 

298 

299 if not isinstance(data, dict): 

300 raise DatasetIntegrityError( 

301 f"Cached dataset {dataset_id!r} has invalid {path.name}: expected a JSON object.", 

302 ) 

303 

304 return data 

305 

306 @staticmethod 

307 def _validate_manifest(manifest: dict[str, object], dataset_id: str) -> None: 

308 required = {"cache_format_version", "dataset_id", "bar_count", "bars_sha256"} 

309 

310 missing = sorted(required - manifest.keys()) 

311 

312 if missing: 

313 raise DatasetIntegrityError( 

314 f"Cached dataset {dataset_id!r} has an invalid " 

315 f"manifest: missing {", ".join(missing)}.", 

316 ) 

317 

318 if manifest["cache_format_version"] != CACHE_FORMAT_VERSION: 

319 raise DatasetIntegrityError( 

320 f"Cached dataset {dataset_id!r} uses unsupported cache " 

321 f"format version {manifest["cache_format_version"]!r}.", 

322 ) 

323 

324 if manifest["dataset_id"] != dataset_id: 

325 raise DatasetIntegrityError( 

326 f"Cached dataset {dataset_id!r} has a manifest for {manifest["dataset_id"]!r}.", 

327 ) 

328 

329 if not isinstance(manifest["bar_count"], int) or isinstance(manifest["bar_count"], bool): 

330 raise DatasetIntegrityError( 

331 f"Cached dataset {dataset_id!r} has an invalid manifest bar_count.", 

332 ) 

333 

334 if not isinstance(manifest["bars_sha256"], str): 

335 raise DatasetIntegrityError( 

336 f"Cached dataset {dataset_id!r} has an invalid manifest bars_sha256.", 

337 )