Python TimedCache

test_timed_cache.py

from time import sleep

from timed_cache import Memoize


def test_memoize():
    """
    Test if the cached function results expires in the expected time.
    """

    @Memoize(ttl=1)
    def mem_sum(a, b):
        return a + b

    mem_sum(1, 1)
    mem_sum(1, 2)

    memoize = mem_sum.__closure__[1].cell_contents
    cache = memoize._cache
    assert not cache[memoize.generate_unique_key(1, 1)].expired()
    assert not cache[memoize.generate_unique_key(1, 2)].expired()

    sleep(1)

    assert cache[memoize.generate_unique_key(1, 1)].expired()
    assert cache[memoize.generate_unique_key(1, 2)].expired()


def test_memoize_in_class():
    class MemoizeTestingClass:
        @Memoize(ttl=1)
        def run(self, a, b):
            return a + b

    executor = MemoizeTestingClass()
    executor.run(1, 1)
    executor.run(1, 2)

    memoize = executor.run.__closure__[1].cell_contents
    cache = memoize._cache
    assert not cache[memoize.generate_unique_key(executor, 1, 1)].expired()
    assert not cache[memoize.generate_unique_key(executor, 1, 2)].expired()

    sleep(1)

    assert cache[memoize.generate_unique_key(executor, 1, 1)].expired()
    assert cache[memoize.generate_unique_key(executor, 1, 2)].expired()


def test_memoize_in_class_with_kwargs():
    class MemoizeTestingClass:
        @Memoize(ttl=1)
        def run(self, a, b):
            return a + b['value']

    executor = MemoizeTestingClass()
    executor.run(1, b={'value': 1})
    executor.run(1, b={'value': 2})

    memoize = executor.run.__closure__[1].cell_contents
    cache = memoize._cache
    assert not cache[memoize.generate_unique_key(executor, 1, b={'value': 1})].expired()
    assert not cache[memoize.generate_unique_key(executor, 1, b={'value': 2})].expired()

    sleep(1)

    assert cache[memoize.generate_unique_key(executor, 1, b={'value': 1})].expired()
    assert cache[memoize.generate_unique_key(executor, 1, b={'value': 2})].expired()

timed_cache.py

"""
This is a simple in-memory cache decorator for timed-limited caching.

#TODO: I think some size limitation here would be helpful.
#TODO: One extra optimization is to avoid hashing the key too many times on cache misses:
https://github.com/python/cpython/blob/master/Lib/functools.py#L398
"""
import collections
import logging
import hashlib
from time import time
from typing import Any

logger = logging.getLogger(__name__)


class CacheEntry:
    """
    A cache entry that stores its value and the creation timestamp.

    Args:
        value:
        ttl: integer, the time-to-live for the entry in seconds.
    """

    def __init__(self, value: Any, ttl: int = 20):
        self.value = value
        self.expires_at = time() + ttl

    def expired(self):
        """Return True if the entry is expired, False otherwise."""
        return self.expires_at < time()


class Memoize:
    """
    A decorator to cache function calls with return values for a limited timespan.
    Supported arguments:
    * ttl: time to live of each cache entry in seconds
    """
    def __init__(self, ttl):
        self._cache = dict()
        self._ttl = ttl
        self._separator = object()

    @staticmethod
    def generate_unique_key(*args, **kwargs):
        """
        Generates a unique key based on the hashed values of all of the passed
        arguments. This makes a pretty bold assumption that the hash() function
        is deterministic, which is (probably) implementation specific.
        """

        hashable_args = [arg if isinstance(arg, collections.abc.Hashable) else tuple(arg) for arg in args]
        hashed_args = [f'{hash(arg)}' for arg in hashable_args]
        hashable_kwargs = [
            (key, value if isinstance(value, collections.abc.Hashable) else tuple(value)) for key, value in kwargs.items()
        ]
        hashed_kwargs = [f'{hash(item)}' for item in hashable_args + hashable_kwargs]
        # this is md5 hashed again to avoid the key growing too large
        return hashlib.md5(':'.join(hashed_args + hashed_kwargs).encode()).hexdigest()

    def __call__(self, func):
        def memoized_func(*args, **kwargs):
            key = self.generate_unique_key(*args, **kwargs)
            if key in self._cache:
                if not self._cache[key].expired():
                    logger.debug('Fetching values for (%s, %s) from the cache.', args, kwargs)
                    return self._cache[key].value
            result = func(*args, **kwargs)
            self._cache[key] = CacheEntry(result, ttl=self._ttl)
            return result

        return memoized_func


memoize = Memoize