Add entry size limit per cache core

Author: shaypal5

README.rst

@@ -138,6 +138,7 @@ The following parameters will only be applied to decorators defined after `set_d
 *  `cache_dir`
 *  `pickle_reload`
 *  `separate_files`
+*  `entry_size_limit`
 
 These parameters can be changed at any time and they will apply to all decorators:
 
@@ -269,6 +270,22 @@ You can specify a maximum allowed age for a cached value on a per-call basis usi
 - If the cached value is older than this threshold, a new calculation is triggered and the cache is updated.
 - If not, the cached value is returned as usual.
 
+Entry Size Limit
+~~~~~~~~~~~~~~~~
+You can prevent very large return values from being cached by specifying
+``entry_size_limit`` on the decorator. Values larger than this limit are
+returned but not stored. The limit accepts an integer number of bytes or a
+human readable string like ``"200MB"``.
+
+.. code-block:: python
+
+  @cachier(entry_size_limit="10KB")
+  def load_data():
+      ...
+
+When ``cachier__verbose=True`` is passed to a call that returns a value
+exceeding the limit, an informative message is printed.
+
 Ignore Cache
 ~~~~~~~~~~~~
 

src/cachier/config.py

@@ -1,6 +1,7 @@
 import hashlib
 import os
 import pickle
+import re
 import threading
 from dataclasses import dataclass, field, replace
 from datetime import datetime, timedelta
@@ -32,6 +33,27 @@ def _default_cache_dir():
     return os.path.expanduser("~/.cachier/")
 
 
+def parse_bytes(size: Union[int, str, None]) -> Optional[int]:
+    """Convert a human friendly size string to bytes."""
+    if size is None:
+        return None
+    if isinstance(size, int):
+        return size
+    match = re.fullmatch(r"(?i)\s*(\d+(?:\.\d+)?)\s*([kmgt]?b)?\s*", str(size))
+    if not match:
+        raise ValueError(f"Invalid size value: {size}")
+    number = float(match.group(1))
+    unit = (match.group(2) or "b").upper()
+    factor = {
+        "B": 1,
+        "KB": 1024,
+        "MB": 1024**2,
+        "GB": 1024**3,
+        "TB": 1024**4,
+    }[unit]
+    return int(number * factor)
+
+
 class LazyCacheDir:
     """Lazily resolve the default cache directory using $XDG_CACHE_HOME."""
 
@@ -65,6 +87,7 @@ class Params:
     allow_none: bool = False
     cleanup_stale: bool = False
     cleanup_interval: timedelta = timedelta(days=1)
+    entry_size_limit: Optional[int] = None
 
 
 _global_params = Params()

src/cachier/core.py

@@ -24,6 +24,7 @@
     HashFunc,
     Mongetter,
     _update_with_defaults,
+    parse_bytes,
 )
 from .cores.base import RecalculationNeeded, _BaseCore
 from .cores.memory import _MemoryCore
@@ -60,11 +61,15 @@ def _function_thread(core, key, func, args, kwds):
         print(f"Function call failed with the following exception:\n{exc}")
 
 
-def _calc_entry(core, key, func, args, kwds) -> Optional[Any]:
+def _calc_entry(
+    core, key, func, args, kwds, printer=lambda *_: None
+) -> Optional[Any]:
     core.mark_entry_being_calculated(key)
     try:
         func_res = func(*args, **kwds)
-        core.set_entry(key, func_res)
+        stored = core.set_entry(key, func_res)
+        if not stored:
+            printer("Result exceeds entry_size_limit; not cached")
         return func_res
     finally:
         core.mark_entry_not_calculated(key)
@@ -123,6 +128,7 @@ def cachier(
     allow_none: Optional[bool] = None,
     cleanup_stale: Optional[bool] = None,
     cleanup_interval: Optional[timedelta] = None,
+    entry_size_limit: Optional[Union[int, str]] = None,
 ):
     """Wrap as a persistent, stale-free memoization decorator.
 
@@ -191,6 +197,10 @@ def cachier(
         thread. Defaults to False.
     cleanup_interval: datetime.timedelta, optional
         Minimum time between automatic cleanup runs. Defaults to one day.
+    entry_size_limit: int or str, optional
+        Maximum serialized size of a cached value. Values exceeding the limit
+        are returned but not cached. Human readable strings like ``"10MB"`` are
+        allowed.
 
     """
     # Check for deprecated parameters
@@ -204,6 +214,9 @@ def cachier(
     # Update parameters with defaults if input is None
     backend = _update_with_defaults(backend, "backend")
     mongetter = _update_with_defaults(mongetter, "mongetter")
+    size_limit_bytes = parse_bytes(
+        _update_with_defaults(entry_size_limit, "entry_size_limit")
+    )
     # Override the backend parameter if a mongetter is provided.
     if callable(mongetter):
         backend = "mongo"
@@ -215,28 +228,34 @@ def cachier(
             cache_dir=cache_dir,
             separate_files=separate_files,
             wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=size_limit_bytes,
         )
     elif backend == "mongo":
         core = _MongoCore(
             hash_func=hash_func,
             mongetter=mongetter,
             wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=size_limit_bytes,
         )
     elif backend == "memory":
         core = _MemoryCore(
-            hash_func=hash_func, wait_for_calc_timeout=wait_for_calc_timeout
+            hash_func=hash_func,
+            wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=size_limit_bytes,
         )
     elif backend == "sql":
         core = _SQLCore(
             hash_func=hash_func,
             sql_engine=sql_engine,
             wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=size_limit_bytes,
         )
     elif backend == "redis":
         core = _RedisCore(
             hash_func=hash_func,
             redis_client=redis_client,
             wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=size_limit_bytes,
         )
     else:
         raise ValueError("specified an invalid core: %s" % backend)
@@ -324,12 +343,12 @@ def _call(*args, max_age: Optional[timedelta] = None, **kwds):
                 )
             key, entry = core.get_entry((), kwargs)
             if overwrite_cache:
-                return _calc_entry(core, key, func, args, kwds)
+                return _calc_entry(core, key, func, args, kwds, _print)
             if entry is None or (
                 not entry._completed and not entry._processing
             ):
                 _print("No entry found. No current calc. Calling like a boss.")
-                return _calc_entry(core, key, func, args, kwds)
+                return _calc_entry(core, key, func, args, kwds, _print)
             _print("Entry found.")
             if _allow_none or entry.value is not None:
                 _print("Cached result found.")
@@ -362,7 +381,7 @@ def _call(*args, max_age: Optional[timedelta] = None, **kwds):
                     try:
                         return core.wait_on_entry_calc(key)
                     except RecalculationNeeded:
-                        return _calc_entry(core, key, func, args, kwds)
+                        return _calc_entry(core, key, func, args, kwds, _print)
                 if _next_time:
                     _print("Async calc and return stale")
                     core.mark_entry_being_calculated(key)
@@ -374,15 +393,15 @@ def _call(*args, max_age: Optional[timedelta] = None, **kwds):
                         core.mark_entry_not_calculated(key)
                     return entry.value
                 _print("Calling decorated function and waiting")
-                return _calc_entry(core, key, func, args, kwds)
+                return _calc_entry(core, key, func, args, kwds, _print)
             if entry._processing:
                 _print("No value but being calculated. Waiting.")
                 try:
                     return core.wait_on_entry_calc(key)
                 except RecalculationNeeded:
-                    return _calc_entry(core, key, func, args, kwds)
+                    return _calc_entry(core, key, func, args, kwds, _print)
             _print("No entry found. No current calc. Calling like a boss.")
-            return _calc_entry(core, key, func, args, kwds)
+            return _calc_entry(core, key, func, args, kwds, _print)
 
         # MAINTAINER NOTE: The main function wrapper is now a standard function
         # that passes *args and **kwargs to _call. This ensures that user

src/cachier/cores/base.py

@@ -9,9 +9,11 @@
 
 import abc  # for the _BaseCore abstract base class
 import inspect
+import pickle
+import sys
 import threading
 from datetime import timedelta
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple
 
 from .._types import HashFunc
 from ..config import CacheEntry, _update_with_defaults
@@ -34,10 +36,12 @@ def __init__(
         self,
         hash_func: Optional[HashFunc],
         wait_for_calc_timeout: Optional[int],
+        entry_size_limit: Optional[int] = None,
     ):
         self.hash_func = _update_with_defaults(hash_func, "hash_func")
         self.wait_for_calc_timeout = wait_for_calc_timeout
         self.lock = threading.RLock()
+        self.entry_size_limit = entry_size_limit
 
     def set_func(self, func):
         """Set the function this core will use.
@@ -90,8 +94,22 @@ def get_entry_by_key(self, key: str) -> Tuple[str, Optional[CacheEntry]]:
 
         """
 
+    def _estimate_size(self, value: Any) -> int:
+        try:
+            return len(pickle.dumps(value))
+        except Exception:
+            return sys.getsizeof(value)
+
+    def _should_store(self, value: Any) -> bool:
+        if self.entry_size_limit is None:
+            return True
+        try:
+            return self._estimate_size(value) <= self.entry_size_limit
+        except Exception:
+            return True
+
     @abc.abstractmethod
-    def set_entry(self, key: str, func_res):
+    def set_entry(self, key: str, func_res: Any) -> bool:
         """Map the given result to the given key in this core's cache."""
 
     @abc.abstractmethod

src/cachier/cores/memory.py

@@ -16,8 +16,9 @@ def __init__(
         self,
         hash_func: Optional[HashFunc],
         wait_for_calc_timeout: Optional[int],
+        entry_size_limit: Optional[int] = None,
     ):
-        super().__init__(hash_func, wait_for_calc_timeout)
+        super().__init__(hash_func, wait_for_calc_timeout, entry_size_limit)
         self.cache: Dict[str, CacheEntry] = {}
 
     def _hash_func_key(self, key: str) -> str:
@@ -29,7 +30,9 @@ def get_entry_by_key(
         with self.lock:
             return key, self.cache.get(self._hash_func_key(key), None)
 
-    def set_entry(self, key: str, func_res: Any) -> None:
+    def set_entry(self, key: str, func_res: Any) -> bool:
+        if not self._should_store(func_res):
+            return False
         hash_key = self._hash_func_key(key)
         with self.lock:
             try:
@@ -47,6 +50,7 @@ def set_entry(self, key: str, func_res: Any) -> None:
                 _condition=cond,
                 _completed=True,
             )
+        return True
 
     def mark_entry_being_calculated(self, key: str) -> None:
         with self.lock:

src/cachier/cores/mongo.py

@@ -40,6 +40,7 @@ def __init__(
         hash_func: Optional[HashFunc],
         mongetter: Optional[Mongetter],
         wait_for_calc_timeout: Optional[int],
+        entry_size_limit: Optional[int] = None,
     ):
         if "pymongo" not in sys.modules:
             warnings.warn(
@@ -49,7 +50,9 @@ def __init__(
             )  # pragma: no cover
 
         super().__init__(
-            hash_func=hash_func, wait_for_calc_timeout=wait_for_calc_timeout
+            hash_func=hash_func,
+            wait_for_calc_timeout=wait_for_calc_timeout,
+            entry_size_limit=entry_size_limit,
         )
         if mongetter is None:
             raise MissingMongetter(
@@ -87,7 +90,9 @@ def get_entry_by_key(self, key: str) -> Tuple[str, Optional[CacheEntry]]:
         )
         return key, entry
 
-    def set_entry(self, key: str, func_res: Any) -> None:
+    def set_entry(self, key: str, func_res: Any) -> bool:
+        if not self._should_store(func_res):
+            return False
         thebytes = pickle.dumps(func_res)
         self.mongo_collection.update_one(
             filter={"func": self._func_str, "key": key},
@@ -104,6 +109,7 @@ def set_entry(self, key: str, func_res: Any) -> None:
             },
             upsert=True,
         )
+        return True
 
     def mark_entry_being_calculated(self, key: str) -> None:
         self.mongo_collection.update_one(

src/cachier/cores/pickle.py

@@ -78,8 +78,9 @@ def __init__(
         cache_dir: Optional[Union[str, os.PathLike]],
         separate_files: Optional[bool],
         wait_for_calc_timeout: Optional[int],
+        entry_size_limit: Optional[int] = None,
     ):
-        super().__init__(hash_func, wait_for_calc_timeout)
+        super().__init__(hash_func, wait_for_calc_timeout, entry_size_limit)
         self._cache_dict: Dict[str, CacheEntry] = {}
         self.reload = _update_with_defaults(pickle_reload, "pickle_reload")
         self.cache_dir = os.path.expanduser(
@@ -119,7 +120,7 @@ def _convert_legacy_cache_entry(
     def _load_cache_dict(self) -> Dict[str, CacheEntry]:
         try:
             with portalocker.Lock(self.cache_fpath, mode="rb") as cf:
-                cache = pickle.load(cf)
+                cache = pickle.load(cf)  # type: ignore[arg-type]
             self._cache_used_fpath = str(self.cache_fpath)
         except (FileNotFoundError, EOFError):
             cache = {}
@@ -146,7 +147,7 @@ def _load_cache_by_key(
         fpath += f"_{hash_str or key}"
         try:
             with portalocker.Lock(fpath, mode="rb") as cache_file:
-                entry = pickle.load(cache_file)
+                entry = pickle.load(cache_file)  # type: ignore[arg-type]
             return _PickleCore._convert_legacy_cache_entry(entry)
         except (FileNotFoundError, EOFError):
             return None
@@ -185,7 +186,7 @@ def _save_cache(
             fpath += f"_{hash_str}"
         with self.lock:
             with portalocker.Lock(fpath, mode="wb") as cf:
-                pickle.dump(cache, cf, protocol=4)
+                pickle.dump(cache, cf, protocol=4)  # type: ignore[arg-type]
             # the same as check for separate_file, but changed for typing
             if isinstance(cache, dict):
                 self._cache_dict = cache
@@ -198,7 +199,9 @@ def get_entry_by_key(
             return key, self._load_cache_by_key(key)
         return key, self.get_cache_dict(reload).get(key)
 
-    def set_entry(self, key: str, func_res: Any) -> None:
+    def set_entry(self, key: str, func_res: Any) -> bool:
+        if not self._should_store(func_res):
+            return False
         key_data = CacheEntry(
             value=func_res,
             time=datetime.now(),
@@ -208,12 +211,13 @@ def set_entry(self, key: str, func_res: Any) -> None:
         )
         if self.separate_files:
             self._save_cache(key_data, key)
-            return  # pragma: no cover
+            return True  # pragma: no cover
 
         with self.lock:
             cache = self.get_cache_dict()
             cache[key] = key_data
             self._save_cache(cache)
+        return True
 
     def mark_entry_being_calculated_separate_files(self, key: str) -> None:
         self._save_cache(