communex.util.memo

  1import time
  2from collections import OrderedDict
  3from collections.abc import MutableMapping
  4from threading import Lock
  5from typing import Callable, Generic, Iterator, TypeVar
  6
  7K = TypeVar("K")
  8V = TypeVar("V")
  9
 10
 11class TTLDict(Generic[K, V], MutableMapping[K, V]):
 12    """
 13    A dictionary that expires its values after a given timeout.
 14
 15    Time tracking is done with `time.monotonic_ns()`.
 16
 17    Be careful, this doens't have and automatic cleanup of expired values yet.
 18    """
 19
 20    ttl: int
 21    _values: OrderedDict[K, tuple[int, V]]
 22    _lock: Lock
 23
 24    def __init__(self, ttl: int, _dict_type: type[OrderedDict[K, tuple[int, V]]] = OrderedDict):
 25        """
 26        Args:
 27            ttl: The timeout in seconds for the memoization.
 28        """
 29        self.ttl = ttl
 30        self._values = _dict_type()
 31        self._lock = Lock()
 32
 33    @property
 34    def ttl_in_ns(self) -> int:
 35        return self.ttl * 1_000_000_000
 36
 37    def __repr__(self) -> str:
 38        return f"<TTLDict@{id(self):#08x} ttl={self.ttl}>"
 39
 40    def __is_expired(self, key: K) -> bool:
 41        expire_time, _ = self._values[key]
 42        return time.monotonic_ns() > expire_time
 43
 44    def __remove_if_expired(self, key: K) -> bool:
 45        expired = self.__is_expired(key)
 46        if expired:
 47            del self._values[key]
 48        return expired
 49
 50    def clean(self):
 51        with self._lock:
 52            for key in self._values.keys():
 53                removed = self.__remove_if_expired(key)
 54                if not removed:
 55                    # TODO: test cleanup optimization
 56                    break
 57
 58    def __setitem__(self, key: K, value: V):
 59        with self._lock:
 60            expire_time = time.monotonic_ns() + self.ttl_in_ns
 61            self._values[key] = (expire_time, value)
 62            self._values.move_to_end(key)
 63
 64    def __getitem__(self, key: K) -> V:
 65        with self._lock:
 66            self.__remove_if_expired(key)
 67            value = self._values[key][1]
 68            return value
 69
 70    def __delitem__(self, key: K):
 71        with self._lock:
 72            del self._values[key]
 73
 74    def __iter__(self) -> Iterator[K]:
 75        with self._lock:
 76            for key in self._values.keys():
 77                if not self.__remove_if_expired(key):
 78                    yield key
 79
 80    def __len__(self) -> int:
 81        """
 82        Warning: this triggers a cleanup, and is O(n) in the number of items in
 83        the dict.
 84        """
 85        self.clean()
 86
 87        # TODO: there is a race condition here.
 88        # Because we are not using RLock as we expect this crap to be used with async.
 89        # But I don't care. Be happy enough with with an "aproximate value".
 90
 91        with self._lock:
 92            return len(self._values)
 93
 94    def get_or_insert_lazy(self, key: K, fn: Callable[[], V]) -> V:
 95        """
 96        Gets the value for the given key, or inserts the value returned by the
 97        given function if the key is not present, returning it.
 98        """
 99        if key in self:
100            return self[key]
101        else:
102            self[key] = fn()
103            return self[key]
104
105
106def __test():
107    m: TTLDict[str, int] = TTLDict(1)
108
109    m["a"] = 2
110
111    print(m.get("a", default="missing"))
112    print(m["a"])
113
114    time.sleep(0.5)
115
116    print(m.get("a", default="missing"))
117    print(m["a"])
118
119    time.sleep(1)
120
121    print(m.get("a", default="missing"))
122    try:
123        print(m["a"])
124
125    except KeyError:
126        print("Key is not present :) yay")
127
128    print(len(m))
129
130    print()
131    print()
132
133    _counter = 0
134
135    def generate():
136        nonlocal _counter
137        _counter += 1
138        print(f"-> generated {_counter}")
139        return _counter
140
141    print("FIRST RUN")
142    v = m.get_or_insert_lazy("a", generate)
143    print(v)
144    v = m.get_or_insert_lazy("a", generate)
145    print(v)
146    print()
147
148    time.sleep(1.5)
149
150    print("SECOND RUN")
151    v = m.get_or_insert_lazy("a", generate)
152    print(v)
153    v = m.get_or_insert_lazy("a", generate)
154    print(v)
155
156
157if __name__ == "__main__":
158    __test()