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