File size: 4,527 Bytes
3d83b62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """Cache management for CiteScan."""
import hashlib
import json
from typing import Any, Callable, Optional
from functools import wraps
from cachetools import TTLCache
import threading
from .config import settings
from .logging import get_logger
logger = get_logger(__name__)
class CacheManager:
"""Thread-safe cache manager using TTLCache."""
def __init__(self):
"""Initialize cache manager."""
self._cache: Optional[TTLCache] = None
self._lock = threading.Lock()
self._initialize_cache()
def _initialize_cache(self) -> None:
"""Initialize the cache based on settings."""
if settings.cache_enabled:
self._cache = TTLCache(
maxsize=settings.cache_max_size,
ttl=settings.cache_ttl
)
logger.info(
f"Cache initialized: max_size={settings.cache_max_size}, "
f"ttl={settings.cache_ttl}s"
)
else:
logger.info("Cache disabled")
def _generate_key(self, *args, **kwargs) -> str:
"""Generate a cache key from arguments.
Args:
*args: Positional arguments
**kwargs: Keyword arguments
Returns:
Cache key as hex string
"""
# Create a stable string representation
key_data = {
"args": args,
"kwargs": sorted(kwargs.items())
}
key_str = json.dumps(key_data, sort_keys=True, default=str)
return hashlib.md5(key_str.encode()).hexdigest()
def get(self, key: str) -> Optional[Any]:
"""Get value from cache.
Args:
key: Cache key
Returns:
Cached value or None if not found
"""
if not settings.cache_enabled or self._cache is None:
return None
with self._lock:
value = self._cache.get(key)
if value is not None:
logger.debug(f"Cache hit: {key}")
return value
def set(self, key: str, value: Any) -> None:
"""Set value in cache.
Args:
key: Cache key
value: Value to cache
"""
if not settings.cache_enabled or self._cache is None:
return
with self._lock:
self._cache[key] = value
logger.debug(f"Cache set: {key}")
def delete(self, key: str) -> None:
"""Delete value from cache.
Args:
key: Cache key
"""
if not settings.cache_enabled or self._cache is None:
return
with self._lock:
if key in self._cache:
del self._cache[key]
logger.debug(f"Cache deleted: {key}")
def clear(self) -> None:
"""Clear all cache entries."""
if not settings.cache_enabled or self._cache is None:
return
with self._lock:
self._cache.clear()
logger.info("Cache cleared")
def get_stats(self) -> dict[str, Any]:
"""Get cache statistics.
Returns:
Dictionary with cache stats
"""
if not settings.cache_enabled or self._cache is None:
return {
"enabled": False,
"size": 0,
"max_size": 0,
"ttl": 0
}
with self._lock:
return {
"enabled": True,
"size": len(self._cache),
"max_size": self._cache.maxsize,
"ttl": self._cache.ttl
}
def cached(self, key_prefix: str = "") -> Callable:
"""Decorator to cache function results.
Args:
key_prefix: Optional prefix for cache key
Returns:
Decorator function
"""
def decorator(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
# Generate cache key
cache_key = f"{key_prefix}:{func.__name__}:{self._generate_key(*args, **kwargs)}"
# Try to get from cache
cached_value = self.get(cache_key)
if cached_value is not None:
return cached_value
# Call function and cache result
result = func(*args, **kwargs)
self.set(cache_key, result)
return result
return wrapper
return decorator
# Global cache manager instance
cache_manager = CacheManager()
|