LRU

lru_sample.py

import heapq
import time
from collections.abc import Callable
from typing import Any


class _Cache:
    def __init__(self, my_fun: Callable, cache_limit: int = 3) -> None:
        self.cache_limit = cache_limit
        self.my_fun = my_fun
        self.records: dict = {}
        self.age_heap: list = []

    def __call__(self, key: str) -> Any:
        if key in self.records:
            self.records[key]["timestamp"] = time.time()
            return self.records[key]["value"]
        else:
            new_record = {"value": self.my_fun(key), "timestamp": time.time()}
            while len(self.age_heap) >= self.cache_limit:
                old_timestamp, old_key = heapq.heappop(self.age_heap)
                old_record = self.records[old_key]
                if old_record["timestamp"] > old_timestamp:
                    print(f"updating timestamp for {old_record} {old_key}")
                    heapq.heappush(
                        self.age_heap, (old_record["timestamp"], old_key)
                    )
                else:
                    del self.records[old_key]
            heapq.heappush(self.age_heap, (new_record["timestamp"], key))
            self.records[key] = new_record
            return new_record["value"]


def Cache(function=None, cache_limit: int = 3):
    if function:
        return _Cache(function)
    else:

        def wrapper(function):
            return _Cache(function, cache_limit)

        return wrapper


def test_cache() -> None:
    @Cache(cache_limit=3)
    def long_running_fun(key: str) -> str:
        return f"###{key}###"

    print(long_running_fun("1"))
    print(long_running_fun("1"))
    print(long_running_fun("2"))
    print(long_running_fun("3"))
    print(long_running_fun.age_heap)
    print(long_running_fun("1"))
    print(long_running_fun.age_heap)
    print(long_running_fun("4"))
    print(long_running_fun.age_heap)
    print(long_running_fun("1"))
    print(long_running_fun.age_heap)
    print(long_running_fun("5"))
    assert long_running_fun("6") == "###6###"
    print(long_running_fun.records.keys())
    for key in long_running_fun.records.keys():
        assert key in ["1", "5", "6"]