ttl_cache.py

from logging import getLogger
from functools import wraps, lru_cache
from typing import Callable, Any
from datetime import datetime

log = getLogger(__name__)


def ttl_cache(
    func: Any = None,
    *,
    initial_ttl: int = 300,
    update_ttl: Callable[[Any], int] = None,
    maxsize: int = 128
):
    """
    Decorator to cache function results with a time-to-live (TTL).
    
    This decorator uses `functools.lru_cache` to store results in memory. The cache is invalidated based on the initial TTL or a dynamically updated TTL.
    
    Args:
        func (Callable): Function to be decorated.
        initial_ttl (int, optional): Initial TTL in seconds. Defaults to 300.
        update_ttl (Callable[[Any], int], optional): Function to update the TTL based on the function's return value. Defaults to None.
        maxsize (int, optional): Maximum size of the cache. Defaults to 128.
    
    Returns:
        Callable: Decorated function with TTL-based caching.
    """
    def decorator(func):
        # Inner function to handle TTL key
        def ttl_func(ttl_key, *args, **kwargs):
            log.debug("TTL key: %s", ttl_key)
            return func(*args, **kwargs)

        # Apply lru_cache to the inner function
        cached_func = lru_cache(maxsize=maxsize)(ttl_func)
        # Initialize the TTL
        cached_func._current_ttl = initial_ttl
        cached_func._start_time = datetime.now()

        @wraps(func)
        def wrapper(*args, **kwargs):
            # Get a timestamp in seconds
            timestamp_s = int(
                round((datetime.now() - cached_func._start_time).total_seconds(), 0)
            )
            # Compute TTL key
            ttl_key = timestamp_s // cached_func._current_ttl
            # Call the cached function with the TTL key
            ret = cached_func(ttl_key, *args, **kwargs)
            # Update TTL if the update_ttl function is provided
            if update_ttl:
                new_ttl = update_ttl(ret)
                cached_func._current_ttl = new_ttl
            return ret

        return wrapper

    if func is not None:
        return decorator(func)
    return decorator