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

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

2

3 import math

4 from collections.abc import MutableMapping

5 from typing import Any, Literal

6

7 from cachetools import LFUCache, LRUCache

8

9 from invariant.cacheable import is_cacheable

10 from invariant.store.base import ArtifactStore

11

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

13

14

15 def _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

25 class 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()