beam: TTL-aware LRU cache module with tests

wizard-work/beam/test_ttl_lru_cache.py

@@ -0,0 +1,93 @@
+import time
+import pytest
+
+from ttl_lru_cache import TTLCache, CacheMiss
+
+
+def test_basic_get_set():
+    c = TTLCache(2)
+    c.set("a", 1)
+    c.set("b", 2)
+    assert c.get("a") == 1
+    assert c.get("z", 9) == 9
+    assert "a" in c
+
+
+def test_lru_eviction():
+    c = TTLCache(2)
+    c.set("a", 1)
+    c.set("b", 2)
+    c.get("a")  # make a recently used
+    c.set("c", 3)
+    assert "b" not in c
+    assert c.get("a") == 1
+    assert c.get("c") == 3
+
+
+def test_ttl_expiry():
+    clock = [0.0]
+    c = TTLCache(10, default_ttl=5, clock=lambda: clock[0])
+    c.set("a", 1)
+    assert c.get("a") == 1
+    clock[0] = 6
+    assert c.get("a") is None
+    assert "a" not in c
+    assert len(c) == 0
+
+
+def test_per_entry_ttl_override():
+    clock = [0.0]
+    c = TTLCache(10, default_ttl=100, clock=lambda: clock[0])
+    c.set("short", 1, ttl=2)
+    c.set("long", 2)
+    clock[0] = 3
+    assert "short" not in c
+    assert c.get("long") == 2
+
+
+def test_peek_does_not_update_recency():
+    c = TTLCache(2)
+    c.set("a", 1)
+    c.set("b", 2)
+    c.peek("a")
+    c.set("c", 3)
+    # a was not touched via get, so it should be evicted (b was more recent? actually b inserted last)
+    # After set a,b: order a,b. peek a no change. set c evicts LRU = a.
+    assert "a" not in c
+
+
+def test_get_strict_raises():
+    c = TTLCache(2)
+    with pytest.raises(CacheMiss):
+        c.get_strict("nope")
+
+
+def test_invalid_args():
+    with pytest.raises(ValueError):
+        TTLCache(0)
+    with pytest.raises(ValueError):
+        TTLCache(2, default_ttl=0)
+    c = TTLCache(2)
+    with pytest.raises(ValueError):
+        c.set("x", 1, ttl=0)
+
+
+def test_delete_and_clear():
+    c = TTLCache(2)
+    c.set("a", 1)
+    assert c.delete("a") is True
+    assert c.delete("a") is False
+    c.set("b", 2)
+    c.clear()
+    assert len(c) == 0
+
+
+def test_capacity_eviction_with_expired():
+    clock = [0.0]
+    c = TTLCache(2, clock=lambda: clock[0])
+    c.set("a", 1, ttl=1)
+    c.set("b", 2, ttl=1)
+    clock[0] = 2  # both expired
+    c.set("c", 3)  # should not count expired toward capacity
+    assert len(c) == 1
+    assert c.get("c") == 3

wizard-work/beam/ttl_lru_cache.py

@@ -0,0 +1,141 @@
+"""TTL-aware LRU cache.
+
+A small, dependency-free, thread-safe-ish (single-threaded design) cache
+combining least-recently-used eviction with per-entry time-to-live expiry.
+
+Design notes:
+- Ordering is tracked with an OrderedDict (move_to_end on access).
+- Expiry is checked lazily on get/peek; expired entries are dropped on access.
+- Capacity counts only live (non-expired) entries at insertion time; expired
+  entries are purged opportunistically when capacity is exceeded.
+- A monotonic clock (time.monotonic) is used so wall-clock jumps don't break TTL.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Callable, Generic, Optional, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class CacheMiss(KeyError):
+    """Raised by get_strict when a key is absent or expired."""
+
+
+@dataclass
+class _Entry(Generic[V]):
+    value: V
+    expires_at: float  # monotonic seconds; inf means no expiry
+
+
+class TTLCache(Generic[K, V]):
+    def __init__(
+        self,
+        capacity: int,
+        default_ttl: Optional[float] = None,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if capacity <= 0:
+            raise ValueError("capacity must be positive")
+        if default_ttl is not None and default_ttl <= 0:
+            raise ValueError("default_ttl must be positive or None")
+        self._capacity = capacity
+        self._default_ttl = default_ttl
+        self._clock = clock
+        self._data: "OrderedDict[K, _Entry[V]]" = OrderedDict()
+
+    # -- internals ---------------------------------------------------------
+
+    def _now(self) -> float:
+        return self._clock()
+
+    def _purge_expired(self) -> None:
+        # Drop expired entries from the front/back opportunistically.
+        now = self._now()
+        dead = [k for k, e in self._data.items() if e.expires_at <= now]
+        for k in dead:
+            del self._data[k]
+
+    def _ttl_to_deadline(self, ttl: Optional[float]) -> float:
+        if ttl is None:
+            ttl = self._default_ttl
+        if ttl is None:
+            return float("inf")
+        if ttl <= 0:
+            raise ValueError("ttl must be positive")
+        return self._now() + ttl
+
+    # -- public API --------------------------------------------------------
+
+    def __len__(self) -> int:
+        self._purge_expired()
+        return len(self._data)
+
+    def __contains__(self, key: K) -> bool:
+        entry = self._data.get(key)
+        if entry is None:
+            return False
+        if entry.expires_at <= self._now():
+            del self._data[key]
+            return False
+        return True
+
+    def get(self, key: K, default: Any = None) -> Any:
+        entry = self._data.get(key)
+        if entry is None:
+            return default
+        if entry.expires_at <= self._now():
+            del self._data[key]
+            return default
+        self._data.move_to_end(key)
+        return entry.value
+
+    def get_strict(self, key: K) -> V:
+        v = self.get(key, _MISSING)
+        if v is _MISSING:
+            raise CacheMiss(key)
+        return v  # type: ignore[return-value]
+
+    def set(self, key: K, value: V, ttl: Optional[float] = None) -> None:
+        deadline = self._ttl_to_deadline(ttl)
+        if key in self._data:
+            del self._data[key]
+        self._data[key] = _Entry(value=value, expires_at=deadline)
+        self._data.move_to_end(key)
+        self._evict_if_needed()
+
+    def _evict_if_needed(self) -> None:
+        self._purge_expired()
+        while len(self._data) > self._capacity:
+            self._data.popitem(last=False)
+
+    def delete(self, key: K) -> bool:
+        if key in self._data:
+            del self._data[key]
+            return True
+        return False
+
+    def clear(self) -> None:
+        self._data.clear()
+
+    def keys(self):
+        self._purge_expired()
+        return list(self._data.keys())
+
+    def items(self):
+        self._purge_expired()
+        return [(k, e.value) for k, e in self._data.items()]
+
+    def peek(self, key: K, default: Any = None) -> Any:
+        """Return value without updating recency or triggering purge of others."""
+        entry = self._data.get(key)
+        if entry is None or entry.expires_at <= self._now():
+            return default
+        return entry.value
+
+
+_MISSING = object()