toolchemy 0.2.185__py3-none-any.whl

toolchemy/utils/at_exit_collector.py

@@ -0,0 +1,109 @@
+from abc import ABC, abstractmethod
+import atexit
+import collections
+
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.utils import pp
+
+
+class ICollectable(ABC):
+    FIELD_NAME_KEY = "name"
+
+    @abstractmethod
+    def label(self) -> str:
+        pass
+
+    @abstractmethod
+    def collect(self) -> dict:
+        pass
+
+
+class AtExitCollector:
+    _collectables: list[ICollectable] = []
+    _collector_logger = get_logger()
+    _is_enabled = False
+    _is_registered = False
+
+    @classmethod
+    def enable(cls):
+        cls._is_enabled = True
+
+    @classmethod
+    def disable(cls):
+        cls._is_enabled = False
+
+    @classmethod
+    def reset(cls):
+        cls._collectables = []
+        atexit.unregister(cls._collector_summary)
+        cls._is_registered = False
+
+    @classmethod
+    def register(cls, collectable: ICollectable):
+        if not cls._is_enabled:
+            return
+        if not cls._is_registered:
+            atexit.register(cls._collector_summary)
+            cls._collector_logger.info("AtExitCollector registered.")
+            cls._is_registered = True
+        cls._collector_logger.info(f"Registering collectable: {collectable} (type: {type(collectable)})")
+        assert isinstance(collectable, ICollectable), f"Expected ICollectable, got {type(collectable)}"
+        cls._collectables.append(collectable)
+        for c in cls._collectables:
+            cls._collector_logger.info(f"Registered collectable: {c} (type: {c})")
+
+    @classmethod
+    def _collector_summary(cls) -> None:
+        if not cls._is_enabled:
+            return
+        if len(cls._collectables) == 0:
+            cls._collector_logger.info(f"No collectable registered, skipping AtExitCollector summary.")
+            return
+        cls._collector_logger.info("AtExitCollector| generating summary...")
+        aggregated = {}
+        for collectable in cls._collectables:
+            try:
+                data = collectable.collect()
+            except TypeError as e:
+                cls._collector_logger.error(f"Collectable of wrong type: {type(collectable)} (is ICollectable: {isinstance(collectable, ICollectable)}): {collectable} (err msg: {e})")
+                raise e
+
+            name = collectable.label()
+            if name not in aggregated:
+                aggregated[name] = {
+                    "instances": 0,
+                }
+            for k, v in data.items():
+                if not isinstance(v, (int, float)):
+                    continue
+                if k not in aggregated[name]:
+                    aggregated[name][k] = 0
+                aggregated[name][k] += v
+            aggregated[name]["instances"] += 1
+
+            cls._collector_logger.info(f"AtExitCollector| summary for {name}:\n{pp(data, print_msg=False)}")
+
+        averages = {}
+        for instance_name, instance_data in aggregated.items():
+            if instance_name not in averages:
+                averages[instance_name] = {}
+            for k, v in instance_data.items():
+                if k == "instances":
+                    continue
+
+                avg_key = f"avg_{k}"
+                averages[instance_name][avg_key] = v / instance_data["instances"]
+
+        cls._update_dict(aggregated, averages)
+
+        cls._collector_logger.info(f"AtExitCollector| aggregated summary:\n{pp(aggregated, print_msg=False)}")
+        cls._collector_logger.info("AtExitCollector| summary generation DONE")
+
+    @classmethod
+    def _update_dict(cls, d, u):
+        for k, v in u.items():
+            if isinstance(v, collections.abc.Mapping):
+                d[k] = cls._update_dict(d.get(k, {}), v)
+            else:
+                d[k] = v
+        return d

toolchemy/utils/cacher/__init__.py

@@ -0,0 +1,20 @@
+from .common import ICacher, BaseCacher, CacheEntryDoesNotExistError, DummyCacher
+from .cacher_pickle import CacherPickle
+from .cacher_shelve import CacherShelve
+from .cacher_diskcache import CacherDiskcache
+
+
+class Cacher(CacherDiskcache):
+    pass
+
+
+__all__ = [
+    "ICacher",
+    "BaseCacher",
+    "Cacher",
+    "CacherPickle",
+    "CacherShelve",
+    "CacherDiskcache",
+    "DummyCacher",
+    "CacheEntryDoesNotExistError",
+]

toolchemy/utils/cacher/cacher_diskcache.py

@@ -0,0 +1,121 @@
+import logging
+import os
+import sys
+from typing import Optional, Any
+from diskcache import Cache
+
+from toolchemy.utils.cacher.common import BaseCacher, CacheEntryDoesNotExistError, CacheEntryHasNotBeenSetError, CacherInitializationError
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.locations import get_external_caller_path
+from toolchemy.utils.utils import _caller_module_name
+
+
+class CacherDiskcache(BaseCacher):
+    def __init__(self, name: str | None = None, cache_base_dir: Optional[str] = None, disabled: bool = False,
+                 log_level: int = logging.INFO):
+        super().__init__()
+        self._disabled = disabled
+
+        self._log_level = log_level
+        self._logger = get_logger(level=self._log_level)
+
+        self._name = name
+        if not self._name:
+            self._name = _caller_module_name()
+
+        self._cache_base_dir = cache_base_dir
+        if self._cache_base_dir is None:
+            self._cache_base_dir = get_external_caller_path()
+
+        self._cache_dir = os.path.join(self._cache_base_dir, self.CACHER_MAIN_NAME, self._name)
+
+        if self._disabled:
+            return
+
+        try:
+            self._cache = Cache(self._cache_dir, cull_limit=0, size_limit=2**38)
+        except Exception as e:
+            raise CacherInitializationError(f"Failed to initialize disk cache for name '{self._name}' (cache dir: '{self._cache_dir}')") from e
+
+        self._logger.debug(
+            f"Cacher '{self._name}' initialized (cache path: '{self._cache_dir}', log_level: '{logging.getLevelName(log_level)}')")
+        self._logger.debug(f"Cacher logging DEBUG level enabled")
+
+    @property
+    def cache_location(self) -> str:
+        return self._cache_dir
+
+    def sub_cacher(self, log_level: int | None = None, suffix: str | None = None) -> "ICacher":
+        name = _caller_module_name()
+        if suffix:
+            name += f"__{suffix}"
+        if log_level is None:
+            log_level = self._log_level
+        self._logger.debug(f"Creating sub cacher")
+        self._logger.debug(f"> base cache dir: {self._cache_dir}")
+        self._logger.debug(f"> name: {name}")
+        self._logger.debug(f"> log level: {log_level} ({logging.getLevelName(log_level)})")
+        self._logger.debug(f"> is disabled: {self._disabled})")
+
+        return CacherDiskcache(name=os.path.join(self._name, name).strip("/"),
+                               cache_base_dir=self._cache_base_dir,
+                               log_level=log_level, disabled=self._disabled)
+
+    def _exists(self, name: str) -> bool:
+        if self._disabled:
+            self._logger.debug("Cacher disabled")
+            return False
+
+        if name in self._cache:
+            return True
+        self._logger.debug("Cache entry %s::%s does not exist", self._cache_dir, name)
+        return False
+
+    def set(self, name: str, content: Any, ttl_s: int | None = None):
+        """
+        Dumps a given object under a given cache entry name. The object must be pickleable.
+        """
+        if self._disabled:
+            return
+
+        result = self._cache.set(name, content, expire=ttl_s)
+        does_exist = self.exists(name)
+        if not result or not does_exist:
+            self._logger.error(f"Cache entry '{name}' not set for name '{self._name}' ({result}, {does_exist})")
+            self._logger.error(f"> cache dir: {self._cache_dir}")
+            self._logger.error(f"> type of the content: {type(content)}")
+            self._logger.error(f"> size of the content: {sys.getsizeof(content)}")
+            raise CacheEntryHasNotBeenSetError()
+
+        self._logger.debug("Cache set %s::%s", self._cache_dir, name)
+
+    def get(self, name: str) -> Any:
+        """
+        Loads an object for a given cache entry name. If it doesn't exist an exception is thrown.
+        """
+        if self._disabled:
+            raise CacheEntryDoesNotExistError(f"Caching is disabled...")
+
+        self._logger.debug("Cache get: %s::%s", self._cache_dir, name)
+
+        if name in self._cache:
+            return self._cache.get(name)
+
+        raise CacheEntryDoesNotExistError(f"Cache does not exist: {self._cache_dir}::{name}.")
+
+    def unset(self, name: str):
+        """
+        Removes a cache entry for a given name
+        """
+        self._logger.debug("Cache unset: %s::%s", self._cache_dir, name)
+        if self._disabled:
+            return
+
+        if name in self._cache:
+            del self._cache[name]
+            self._logger.debug("Cache entry %s::%s removed", self._name, name)
+        else:
+            self._logger.warning("Cache entry %s::%s does not exist, nothing to remove", self._name, name)
+
+    def persist(self):
+        self._cache.close()

toolchemy/utils/cacher/cacher_pickle.py

@@ -0,0 +1,152 @@
+import logging
+import os
+import pickle
+import threading
+from pathlib import Path
+from typing import Optional, Any
+import shutil
+
+from toolchemy.utils.cacher.common import BaseCacher, DummyLock, CacheEntryDoesNotExistError
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.locations import get_external_caller_path
+from toolchemy.utils.utils import _caller_module_name
+from toolchemy.utils.datestimes import current_unix_timestamp
+
+
+class CacherPickle(BaseCacher):
+    """
+    Cacher implementation where cache is stored as a pickled local file
+    """
+
+    CACHER_MAIN_NAME = ".cache"
+
+    def __init__(self, name: str | None = None, cache_base_dir: Optional[str] = None, disabled: bool = False,
+                 log_level: int = logging.INFO, enable_thread_safeness: bool = False):
+        """
+        Initialize cache with its name. It creates .cache/name subdir in cache_base_dir directory
+        """
+        super().__init__()
+        self._disabled = disabled
+        if enable_thread_safeness:
+            self._lock = threading.Lock()
+        else:
+            self._lock = DummyLock()
+        self._log_level = log_level
+        self._logger = get_logger(level=self._log_level)
+
+        self._name = name
+        if not self._name:
+            self._name = _caller_module_name()
+
+        self._cache_base_dir = cache_base_dir
+        if self._cache_base_dir is None:
+            self._cache_base_dir = get_external_caller_path()
+
+        self._cache_dir = os.path.join(self._cache_base_dir, self.CACHER_MAIN_NAME, self._name)
+
+        with self._lock:
+            os.makedirs(self._cache_dir, exist_ok=True)
+
+        self._logger.debug(f"Cacher '{self._name}' initialized (cache dir: '{self._cache_dir}', log_level: '{logging.getLevelName(log_level)}')")
+        self._logger.debug(f"Cacher logging DEBUG level enabled")
+
+    @property
+    def cache_location(self) -> str:
+        return self._cache_dir
+
+    def sub_cacher(self, log_level: int | None = None, suffix: str | None = None) -> "ICacher":
+        name = _caller_module_name()
+        if suffix:
+            name += f"__{suffix}"
+        if log_level is None:
+            log_level = self._log_level
+        self._logger.debug(f"Creating sub cacher")
+        self._logger.debug(f"> base name: {self._name}")
+        self._logger.debug(f"> name: {name}")
+        self._logger.debug(f"> log level: {log_level} ({logging.getLevelName(log_level)})")
+        self._logger.debug(f"> is disabled: {self._disabled})")
+        return CacherPickle(name=os.path.join(self._name, name).strip("/"), cache_base_dir=self._cache_base_dir,
+                            log_level=log_level, disabled=self._disabled)
+
+    def _exists(self, name: str) -> bool:
+        if self._disabled:
+            self._logger.debug("Cacher disabled")
+            return False
+        target_filename = self._cache_name(name)
+        target_file = Path(target_filename)
+
+        ret_val = False
+        if target_file.is_file() or target_file.is_symlink():
+            try:
+                entry = self._get(name)
+                if entry['ttl_s'] is None:
+                    ret_val = True
+                else:
+                    current_time = current_unix_timestamp()
+                    if current_time - entry['timestamp'] < entry['ttl_s']:
+                        ret_val = True
+                    else:
+                        self.unset(name)
+                        ret_val = False
+            except CacheEntryDoesNotExistError:
+                return False
+            self._logger.debug("Cache entry %s::%s (%s) exists", self._name, name, target_filename)
+
+        self._logger.debug("Cache entry %s::%s does not exist", self._name, name)
+        return ret_val
+
+    def set(self, name: str, content: Any, ttl_s: int | None = None):
+        if self._disabled:
+            return
+        target_filename = self._cache_name(name)
+        with self._lock:
+            with open(target_filename, "wb") as file:
+                try:
+                    pickle.dump(self._envelop(content, ttl_s=ttl_s), file)  # type: ignore
+                except TypeError as e:
+                    self._logger.error(f"Wrong type of the serialized content: {type(content)}. Target: {target_filename}. Content:\n{content}")
+                    shutil.rmtree(target_filename)
+                    raise e
+        self._logger.debug("Cache set %s::%s (file: %s)", self._name, name, target_filename)
+
+    def get(self, name: str) -> Any:
+        entry = self._get(name)
+        return entry["data"]
+
+    def _get(self, name: str) -> Any:
+        if self._disabled:
+            raise CacheEntryDoesNotExistError(f"Caching is disabled...")
+        target_filename = self._cache_name(name)
+        self._logger.debug("Cache get: %s::%s (file: %s)", self._name, name, target_filename)
+        target_file = Path(target_filename)
+        with self._lock:
+            if target_file.is_file() or target_file.is_symlink():
+                with open(target_filename, "rb") as file:
+                    try:
+                        restored_object = pickle.load(file, encoding="utf-8")
+                    except ModuleNotFoundError as e:
+                        self._logger.error(f"{e} while loading from file: '{target_filename}'")
+                        raise e
+                if restored_object['ttl_s'] is not None:
+                    current_time = current_unix_timestamp()
+                    if current_time - restored_object['timestamp'] >= restored_object['ttl_s']:
+                        self.unset(name)
+                        raise CacheEntryDoesNotExistError(f"Cache does not exist: {self._name}::{name}. Path: {target_filename}")
+                return restored_object
+        raise CacheEntryDoesNotExistError(f"Cache does not exist: {self._name}::{name}. Path: {target_filename}")
+
+    def unset(self, name: str):
+        """
+        Removes a cache entry for a given name
+        """
+        target_filename = self._cache_name(name)
+        self._logger.debug("Cache unset: %s::%s (file: %s)", self._name, name, target_filename)
+        target_file = Path(target_filename)
+        if target_file.is_file() or target_file.is_symlink():
+            os.remove(target_filename)
+            self._logger.debug("Cache entry %s::%s removed", self._name, name)
+        else:
+            self._logger.warning("Cache entry %s::%s does not exist, nothing to remove", self._name, name)
+
+    def _cache_name(self, name: str):
+        return os.path.join(self._cache_dir, f"{name}.pkl")

toolchemy/utils/cacher/cacher_shelve.py

@@ -0,0 +1,196 @@
+import logging
+import os
+import threading
+from typing import Optional, Any
+import shelve
+
+from toolchemy.utils.cacher.common import BaseCacher, DummyLock, CacheEntryDoesNotExistError
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.locations import get_external_caller_path
+from toolchemy.utils.utils import _caller_module_name
+from toolchemy.utils.datestimes import current_unix_timestamp
+
+
+class CacherShelve(BaseCacher):
+    def __init__(self, name: str | None = None, cache_base_dir: Optional[str] = None, disabled: bool = False,
+                 log_level: int = logging.INFO, enable_thread_safeness: bool = False):
+        super().__init__()
+        self._disabled = disabled
+        self._enable_thread_safeness = enable_thread_safeness
+        if enable_thread_safeness:
+            self._lock = threading.Lock()
+        else:
+            self._lock = DummyLock()
+        self._log_level = log_level
+        self._logger = get_logger(level=self._log_level)
+
+        self._name = name
+        if not self._name:
+            self._name = _caller_module_name()
+
+        self._cache_base_dir = cache_base_dir
+        if self._cache_base_dir is None:
+            self._cache_base_dir = get_external_caller_path()
+
+        self._cache_dir = os.path.join(self._cache_base_dir, self.CACHER_MAIN_NAME, self._name)
+
+        if self._disabled:
+            return
+
+        self._cache_path = os.path.join(self._cache_dir, "cache")
+
+        with self._lock:
+            os.makedirs(self._cache_dir, exist_ok=True)
+        if not self._enable_thread_safeness:
+            self._open()
+
+        self._logger.debug(f"Cacher '{self._name}' initialized (cache path: '{self._cache_dir}', log_level: '{logging.getLevelName(log_level)}')")
+        self._logger.debug(f"Cacher logging DEBUG level enabled")
+
+    def _open(self):
+        self._cache = shelve.open(self._cache_path, writeback=False)
+
+    @property
+    def cache_location(self) -> str:
+        return self._cache_dir
+
+    def sub_cacher(self, log_level: int | None = None, suffix: str | None = None) -> "ICacher":
+        name = _caller_module_name()
+        if suffix:
+            name += f"__{suffix}"
+        if log_level is None:
+            log_level = self._log_level
+        self._logger.debug(f"Creating sub cacher")
+        self._logger.debug(f"> base cache dir: {self._cache_dir}")
+        self._logger.debug(f"> name: {name}")
+        self._logger.debug(f"> log level: {log_level} ({logging.getLevelName(log_level)})")
+        self._logger.debug(f"> is disabled: {self._disabled})")
+
+        return CacherShelve(name=os.path.join(self._name, name).strip("/"),
+                            cache_base_dir=self._cache_base_dir,
+                            log_level=log_level, disabled=self._disabled)
+
+
+    def _exists(self, name: str) -> bool:
+        if self._disabled:
+            self._logger.debug("Cacher disabled")
+            return False
+
+        ret_val = False
+        with self._lock:
+            if self._enable_thread_safeness:
+                self._open()
+            if name in self._cache:
+                self._logger.debug("Cache entry %s::%s exists", self._cache_dir, name)
+                try:
+                    existing_entry = self._cache[name]
+                except Exception as e:
+                    self._logger.error(f"Cache entry failed to fetch cached entry: {e}")
+                    self._logger.error(f"Existing keys: {self._cache.keys()}")
+                    if self._enable_thread_safeness:
+                        self._close()
+                    raise e
+                self._logger.debug(f"Cache existing entry: {existing_entry}")
+                entry = self._migrate(name, existing_entry)
+                if self._cache[name]['ttl_s'] is None:
+                    ret_val = True
+                else:
+                    current_time = current_unix_timestamp()
+                    if current_time - entry['timestamp'] < entry['ttl_s']:
+                        ret_val = True
+                    else:
+                        del self._cache[name]
+            if self._enable_thread_safeness:
+                self._close()
+            self._logger.debug("Cache entry %s::%s does not exist", self._cache_dir, name)
+        return ret_val
+
+    def set(self, name: str, content: Any, ttl_s: int | None = None):
+        """
+        Dumps a given object under a given cache entry name. The object must be pickleable.
+        """
+        if self._disabled:
+            return
+
+        with self._lock:
+            if self._enable_thread_safeness:
+                self._open()
+            self._cache[name] = self._envelop(content, ttl_s=ttl_s)
+            if self._enable_thread_safeness:
+                self._close()
+
+        self._logger.debug("Cache set %s::%s", self._cache_dir, name)
+
+    def get(self, name: str) -> Any:
+        """
+        Loads an object for a given cache entry name. If it doesn't exist an exception is thrown.
+        """
+        if self._disabled:
+            raise CacheEntryDoesNotExistError(f"Caching is disabled...")
+
+        self._logger.debug("Cache get: %s::%s", self._cache_dir, name)
+
+        ret_val = None
+
+        with self._lock:
+            if self._enable_thread_safeness:
+                self._open()
+            if name in self._cache:
+                entry = self._migrate(name, self._cache[name])
+                if entry['ttl_s'] is None:
+                    ret_val = entry['data']
+                else:
+                    current_time = current_unix_timestamp()
+                    if current_time - entry['timestamp'] < entry['ttl_s']:
+                        ret_val = entry['data']
+                    else:
+                        del self._cache[name]
+            if self._enable_thread_safeness:
+                self._close()
+        if ret_val is not None:
+            return ret_val
+        raise CacheEntryDoesNotExistError(f"Cache does not exist: {self._cache_dir}::{name}.")
+
+    def unset(self, name: str):
+        """
+        Removes a cache entry for a given name
+        """
+        self._logger.debug("Cache unset: %s::%s", self._cache_dir, name)
+        if self._disabled:
+            return
+
+        if name in self._cache:
+            with self._lock:
+                if self._enable_thread_safeness:
+                    self._open()
+                del self._cache[name]
+                if self._enable_thread_safeness:
+                    self._close()
+            self._logger.debug("Cache entry %s::%s removed", self._name, name)
+        else:
+            self._logger.warning("Cache entry %s::%s does not exist, nothing to remove", self._name, name)
+
+    def _close(self):
+        self._cache.close()
+
+    def persist(self):
+        if not self._enable_thread_safeness:
+            with self._lock:
+                self._close()
+
+    def _migrate(self, name: str, entry: Any) -> dict[str, Any]:
+        if not isinstance(entry, dict) or ("data" not in entry and "timestamp" not in entry and "ttl_s" not in entry):
+            self._logger.info(f"Migrating data entry to handle TTL")
+            self._logger.info(f"> entry: {entry} (type: {type(entry)})")
+            self.set(name, entry)
+            entry = self._cache[name]
+        return entry
+
+
+def testing():
+    cacher = CacherShelve()
+    print(cacher._name)
+
+
+if __name__ == "__main__":
+    testing()