Coverage for src / invariant / store / memory.py: 96.30%

1"""MemoryStore: In-memory artifact storage for testing.""" 

2 

3import math 

4from collections.abc import MutableMapping 

5from typing import Any, Literal 

6 

7from cachetools import LFUCache, LRUCache 

8 

9from invariant.cacheable import is_cacheable 

10from invariant.store.base import ArtifactStore 

11 

12CachePolicy = Literal["unbounded", "lru", "lfu"] 

13 

14 

15def _create_cache( 

16 max_size: int, 

17 cache: Literal["lru", "lfu"], 

18) -> MutableMapping[str, Any]: 

19 """Create a cachetools cache.""" 

20 if cache == "lfu": 

21 return LFUCache(maxsize=max_size) 

22 return LRUCache(maxsize=max_size) 

23 

24 

25class MemoryStore(ArtifactStore): 

26 """In-memory artifact store using a dictionary or cachetools cache. 

27 

28 Fast and ephemeral - suitable for testing. Artifacts are lost when 

29 the store instance is destroyed. 

30 

31 Default is cache="lru" with max_size=1000 (safe bounded). Use 

32 cache="unbounded" for explicit unbounded. Use cache="lfu" for 

33 least-frequently-used eviction (e.g. graphics pipelines). 

34 """ 

35 

36 def __init__( 

37 self, 

38 cache: CachePolicy | MutableMapping[str, Any] = "lru", 

39 max_size: int | None = None, 

40 ) -> None: 

41 """Initialize memory store. 

42 

43 Args: 

44 cache: "unbounded" (plain dict), "lru", "lfu", or a MutableMapping 

45 instance (e.g. cachetools.TTLCache). Default "lru". 

46 max_size: Maximum number of artifacts. Default 1000 when cache is 

47 "lru" or "lfu". Ignored for "unbounded". Must be None for 

48 cache instance. 

49 

50 Raises: 

51 ValueError: Invalid combination of cache and max_size. 

52 """ 

53 super().__init__() 

54 

55 # Cache instance provided: use it, max_size forbidden 

56 if isinstance(cache, MutableMapping): 

57 if max_size is not None: 

58 raise ValueError( 

59 "max_size must not be set when cache is a cache instance" 

60 ) 

61 self._artifacts: dict[str, Any] | MutableMapping[str, Any] = cache 

62 return 

63 

64 # Unbounded: plain dict 

65 if cache == "unbounded": 

66 self._artifacts = {} 

67 return 

68 

69 # LRU or LFU: validate max_size, use default 1000 if None 

70 if cache in ("lru", "lfu"): 

71 size = max_size if max_size is not None else 1000 

72 if size < 1: 

73 raise ValueError("max_size must be at least 1") 

74 if size == math.inf or (isinstance(size, float) and math.isinf(size)): 

75 raise ValueError("max_size cannot be infinity") 

76 self._artifacts = _create_cache(size, cache) 

77 return 

78 

79 raise ValueError( 

80 f"cache must be 'unbounded', 'lru', 'lfu', or a MutableMapping; got {cache!r}" 

81 ) 

82 

83 def _make_key(self, op_name: str, digest: str) -> str: 

84 """Create a composite key from op_name and digest.""" 

85 return f"{op_name}:{digest}" 

86 

87 def exists(self, op_name: str, digest: str) -> bool: 

88 """Check if an artifact exists.""" 

89 key = self._make_key(op_name, digest) 

90 exists = key in self._artifacts 

91 if exists: 

92 self.stats.hits += 1 

93 else: 

94 self.stats.misses += 1 

95 return exists 

96 

97 def get(self, op_name: str, digest: str) -> Any: 

98 """Retrieve an artifact by operation name and digest. 

99 

100 Raises: 

101 KeyError: If artifact does not exist. 

102 """ 

103 key = self._make_key(op_name, digest) 

104 if key not in self._artifacts: 

105 raise KeyError( 

106 f"Artifact with op_name '{op_name}' and digest '{digest}' not found" 

107 ) 

108 

109 # Return stored object directly (no deserialization needed) 

110 return self._artifacts[key] 

111 

112 def put(self, op_name: str, digest: str, artifact: Any) -> None: 

113 """Store an artifact with the given operation name and digest.""" 

114 # Validate artifact is cacheable 

115 if not is_cacheable(artifact): 

116 raise TypeError( 

117 f"Artifact is not cacheable: {type(artifact)}. " 

118 f"Use is_cacheable() to check values before storing." 

119 ) 

120 

121 # Store object directly (no serialization needed - relies on immutability contract) 

122 key = self._make_key(op_name, digest) 

123 self._artifacts[key] = artifact 

124 self.stats.puts += 1 

125 

126 def clear(self) -> None: 

127 """Clear all artifacts (mainly for testing).""" 

128 self._artifacts.clear() 

129 self.reset_stats()