Coverage for src / monte_neo / utils / cache.py: 88%

1"""Cache utility for Monte-Neo. 

2 

3Handles caching of compiled Metal libraries, calibration results, and other expensive operations. 

4""" 

5 

6from __future__ import annotations 

7 

8import hashlib 

9import json 

10import os 

11import pickle 

12from typing import Any 

13 

14from monte_neo.utils.logger import get_logger 

15 

16logger = get_logger(__name__) 

17 

18CACHE_DIR = os.path.expanduser("~/.cache/monte_neo") 

19 

20def get_cache_path(filename: str) -> str: 

21 """Get absolute path for a cache file.""" 

22 if not os.path.exists(CACHE_DIR): 

23 os.makedirs(CACHE_DIR, exist_ok=True) 

24 return os.path.join(CACHE_DIR, filename) 

25 

26def save_cache(name: str, data: Any, use_pickle: bool = False) -> bool: 

27 """Save data to cache.""" 

28 try: 

29 path = get_cache_path(name) 

30 if use_pickle: 

31 with open(path, "wb") as f: 

32 pickle.dump(data, f) 

33 else: 

34 with open(path, "w") as f: 

35 json.dump(data, f) 

36 return True 

37 except Exception as e: 

38 logger.warning(f"Failed to save cache {name}: {e}") 

39 return False 

40 

41def load_cache(name: str, use_pickle: bool = False) -> Any | None: 

42 """Load data from cache.""" 

43 path = get_cache_path(name) 

44 if not os.path.exists(path): 

45 return None 

46 

47 try: 

48 if use_pickle: 

49 with open(path, "rb") as f: 

50 return pickle.load(f) 

51 else: 

52 with open(path) as f: 

53 return json.load(f) 

54 except Exception as e: 

55 logger.warning(f"Failed to load cache {name}: {e}") 

56 return None 

57 

58 

59def get_data_hash(data: Any) -> str: 

60 """Generate a hash for data to use as cache key.""" 

61 if hasattr(data, "values"): 

62 # For pandas objects, hash the values 

63 return hashlib.md5(data.values.tobytes()).hexdigest() 

64 return hashlib.md5(str(data).encode()).hexdigest() 

65 

66def save_calibration(indicator_name: str, data: Any, params: dict[str, Any]) -> bool: 

67 """Save optimized parameters to cache.""" 

68 data_hash = get_data_hash(data) 

69 cache_name = f"calibration_{indicator_name}_{data_hash}.json" 

70 return save_cache(cache_name, params) 

71 

72def load_calibration(indicator_name: str, data: Any) -> dict[str, Any] | None: 

73 """Load optimized parameters from cache.""" 

74 data_hash = get_data_hash(data) 

75 cache_name = f"calibration_{indicator_name}_{data_hash}.json" 

76 return load_cache(cache_name) 

77 

78def clear_cache(name: str | None = None) -> None: 

79 """Clear cache files.""" 

80 if name: 

81 path = get_cache_path(name) 

82 if os.path.exists(path): 

83 os.remove(path) 

84 else: 

85 if os.path.exists(CACHE_DIR): 

86 for f in os.listdir(CACHE_DIR): 

87 os.remove(os.path.join(CACHE_DIR, f))