Source code for imgutils.utils.cache

"""
This module provides a thread-safe version of Python's built-in lru_cache decorator.

The main component of this module is the ts_lru_cache decorator, which wraps the standard
lru_cache with thread-safe functionality. This is particularly useful in multi-threaded
environments where cache access needs to be synchronized to prevent race conditions.

Usage:
    >>> from imgutils.utils import ts_lru_cache
    ...
    >>> @ts_lru_cache(maxsize=100)
    >>> def expensive_function(x, y):
    ...     # Some expensive computation
    ...     return x + y
"""
import os
import threading
from collections import defaultdict
from functools import lru_cache, wraps

from typing import Literal

__all__ = ['ts_lru_cache']

LevelTyping = Literal['global', 'process', 'thread']


def _get_context_key(level: LevelTyping = 'global'):
    """
    Get a context key based on the specified caching level.

    :param level: The caching level to use. Can be 'global', 'process', or 'thread'.
    :type level: LevelTyping

    :return: A context key appropriate for the specified level.
    :rtype: tuple or None

    :raises ValueError: If an invalid cache level is specified.

    .. note::
        The function returns:

        - None for 'global' level
        - Process ID for 'process' level
        - (Process ID, Thread ID) tuple for 'thread' level
    """
    if level == 'global':
        return None
    elif level == 'process':
        return os.getpid()
    elif level == 'thread':
        return os.getpid(), threading.get_ident()
    else:
        raise ValueError(f'Invalid cache level, '
                         f'\'global\', \'process\' or \'thread\' expected but {level!r} found.')


[docs]def ts_lru_cache(level: LevelTyping = 'global', **options): """ A thread-safe version of the lru_cache decorator. This decorator wraps the standard lru_cache with a threading lock to ensure thread-safety in multithreaded environments. It maintains the same interface as the built-in lru_cache, allowing you to specify options like maxsize. :param level: The caching level ('global', 'process', or 'thread'). :type level: LevelTyping :param options: Keyword arguments to be passed to the underlying lru_cache. :type options: dict :return: A thread-safe cached version of the decorated function. :rtype: function :example: >>> @ts_lru_cache(level='thread', maxsize=100) >>> def my_function(x, y): ... # Function implementation ... return x + y .. note:: The decorator provides three levels of caching: - global: Single cache shared across all processes and threads - process: Separate cache for each process - thread: Separate cache for each thread .. note:: While this decorator ensures thread-safety, it may introduce some overhead due to lock acquisition. Use it when thread-safety is more critical than maximum performance in multithreaded scenarios. .. note:: The decorator preserves the cache_info() and cache_clear() methods from the original lru_cache implementation. """ _ = _get_context_key(level) def _decorator(func): """ Inner decorator function that wraps the original function. :param func: The function to be decorated. :type func: function :return: The wrapped function with thread-safe caching. :rtype: function """ @lru_cache(**options) @wraps(func) def _cached_func(*args, __context_key=None, **kwargs): """ Cached version of the original function. :param args: Positional arguments to be passed to the original function. :param __context_key: Internal context key for cache separation. :param kwargs: Keyword arguments to be passed to the original function. :return: The result of the original function call. """ return func(*args, **kwargs) lock_pool = defaultdict(threading.Lock) lock = threading.Lock() @wraps(_cached_func) def _new_func(*args, **kwargs): """ Thread-safe wrapper around the cached function. This function acquires a lock before calling the cached function, ensuring thread-safety. :param args: Positional arguments to be passed to the cached function. :param kwargs: Keyword arguments to be passed to the cached function. :return: The result of the cached function call. """ context_key = _get_context_key(level=level) with lock: _context_lock = lock_pool[context_key] with _context_lock: return _cached_func(*args, __context_key=context_key, **kwargs) # Preserve cache_info and cache_clear methods if they exist if hasattr(_cached_func, 'cache_info'): _new_func.cache_info = _cached_func.cache_info if hasattr(_cached_func, 'cache_clear'): _new_func.cache_clear = _cached_func.cache_clear return _new_func return _decorator