yes3 0.0.1__py3-none-any.whl

yes3/__init__.py

@@ -0,0 +1 @@
+from .s3 import S3Location, is_s3_url

yes3/caching/__init__.py

@@ -0,0 +1,6 @@
+from .base import Cache, CacheCore, CachedItemMeta, Serializer, check_meta_mismatches
+from .local_cache import LocalDiskCache
+from .memory_cache import MemoryCache
+from .multi_cache import MultiCache
+from .s3_cache import S3Cache
+from .setup_helpers import setup_cache

yes3/caching/base.py

@@ -0,0 +1,302 @@
+from abc import ABCMeta, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime, UTC
+from typing import Iterable, Iterator, Optional, Self
+
+_NotSpecified = object()
+
+
+def raise_not_found(key) -> KeyError:
+    raise KeyError(f"key '{key}' not found in cache")
+
+
+@dataclass
+class CachedItemMeta:
+    key: str
+    path: Optional[str]
+    size: Optional[int]
+    timestamp: Optional[datetime]
+
+    _ts_format = '%Y-%m-%d %H:%M:%S.%f %z'
+
+    def __post_init__(self):
+        if isinstance(self.timestamp, float):
+            self.timestamp = datetime.fromtimestamp(self.timestamp, UTC)
+        if isinstance(self.timestamp, str):
+            self.timestamp = datetime.strptime(self.timestamp, self._ts_format)
+
+    def to_dict(self) -> dict:
+        return {
+            'key': self.key,
+            'path': self.path,
+            'size': self.size,
+            'timestamp': self.timestamp.strftime(self._ts_format) if self.timestamp else None,
+        }
+
+
+class CacheCore(metaclass=ABCMeta):
+    def __init__(self, active=True, read_only=False):
+        self._read_only = read_only
+        self._active = active
+
+    @abstractmethod
+    def __contains__(self, key):
+        pass
+
+    @abstractmethod
+    def get(self, key, default=_NotSpecified):
+        pass
+
+    @abstractmethod
+    def get_meta(self, key) -> CachedItemMeta:
+        pass
+
+    @abstractmethod
+    def put(self, key, obj, update=False, meta: Optional[CachedItemMeta] = None):
+        pass
+
+    @abstractmethod
+    def remove(self, key):
+        pass
+
+    @abstractmethod
+    def keys(self):
+        pass
+
+    def __getitem__(self, key: str):
+        return self.get(key)
+
+    def __setitem__(self, key: str, obj) -> None:
+        self.put(key, obj)
+
+    def __delitem__(self, key: str) -> None:
+        self.remove(key)
+
+    def is_active(self) -> bool:
+        return self._active
+
+    def activate(self):
+        self._active = True
+        return self
+
+    def deactivate(self):
+        self._active = False
+        return self
+
+    def is_read_only(self) -> bool:
+        return self._read_only
+
+    def set_read_only(self, value: bool) -> Self:
+        self._read_only = value
+        return self
+
+    def update(self, key: str, obj):
+        if key not in self:
+            raise_not_found(key)
+        self.put(key, obj, update=True)
+
+    def pop(self, key: str, default=_NotSpecified):
+        obj = self.get(key, default=default)
+        self.remove(key)
+        return obj
+
+    def list(self) -> dict[str, CachedItemMeta]:
+        items_meta = {}
+        for key in self.keys():
+            items_meta[key] = self.get_meta(key)
+        return items_meta
+
+    def subcache(self, *args, **kwargs) -> Self:
+        raise NotImplementedError(f"`subcache` method is not defined for class {type(self).__name__}")
+
+
+class CacheReaderWriter(metaclass=ABCMeta):
+    @abstractmethod
+    def read(self, key: str):
+        pass
+
+    @abstractmethod
+    def get_meta(self, key: str) -> CachedItemMeta:
+        pass
+
+    @abstractmethod
+    def write(self, key: str, obj, meta=None) -> CachedItemMeta:
+        pass
+
+    @abstractmethod
+    def delete(self, key: str, meta_only=False):
+        pass
+
+
+class CacheCatalog(metaclass=ABCMeta):
+    @abstractmethod
+    def contains(self, key: str):
+        pass
+
+    @abstractmethod
+    def add(self, key: str, info: CachedItemMeta):
+        pass
+
+    @abstractmethod
+    def get(self, key: str) -> CachedItemMeta:
+        pass
+
+    @abstractmethod
+    def remove(self, key: str):
+        pass
+
+    @abstractmethod
+    def keys(self):
+        pass
+
+    @abstractmethod
+    def items(self):
+        pass
+
+
+_CatalogT = dict[str, CachedItemMeta]
+_CatalogBuilderT = Callable[[], _CatalogT]
+
+
+class CacheDictCatalog(CacheCatalog):
+    def __init__(
+            self,
+            catalog: Optional[dict[str, CachedItemMeta]] = None,
+            catalog_builder: Optional[_CatalogBuilderT] = None,
+    ):
+        self._catalog = catalog
+        if catalog_builder is None:
+            catalog_builder = dict
+        self._build_catalog = catalog_builder
+        if self._catalog is None:
+            self.rebuild()
+
+    def rebuild(self):
+        self._catalog = self._build_catalog().copy()
+
+    def contains(self, key: str):
+        return str(key) in self._catalog
+
+    def add(self, key: str, meta: CachedItemMeta):
+        self._catalog[str(key)] = meta
+
+    def get(self, key: str) -> CachedItemMeta:
+        return self._catalog[str(key)]
+
+    def remove(self, key: str):
+        self._catalog.pop(str(key))
+
+    def keys(self):
+        return list(self._catalog.keys())
+
+    def items(self) -> Iterator[tuple[str, CachedItemMeta]]:
+        return iter(self._catalog.items())
+
+
+class Cache(CacheCore, metaclass=ABCMeta):
+    def __init__(self, catalog: CacheCatalog, reader_writer: CacheReaderWriter, active=True, read_only=False):
+        super().__init__(active=active, read_only=read_only)
+        self._catalog = catalog
+        self._reader_writer = reader_writer
+
+    @classmethod
+    @abstractmethod
+    def create(cls, *args, **kwargs):
+        pass
+
+    def __contains__(self, key: str) -> bool:
+        if not self.is_active():
+            return False
+        return self._catalog.contains(key)
+
+    def get(self, key: str, default=_NotSpecified):
+        if not self.is_active() or key not in self:
+            if default is _NotSpecified:
+                raise_not_found(key)
+            else:
+                return default
+        return self._reader_writer.read(key)
+
+    def get_meta(self, key: str) -> CachedItemMeta:
+        if not self.is_active() or key not in self:
+            raise_not_found(key)
+        return self._catalog.get(key)
+
+    def put(self, key: str, obj, *, update=False, meta: Optional[CachedItemMeta] = None) -> Self:
+        if self.is_read_only():
+            raise TypeError('Cache is in read only mode')
+        if self.is_active():
+            if key in self and not update:
+                raise ValueError(f"key '{key}' already exists in cache; use 'update' to overwrite")
+            meta = self._reader_writer.write(key, obj, meta=meta)
+            self._catalog.add(key, meta)
+        else:
+            print(f'WARNING: {type(self).__name__} is not active')
+        return self
+
+    def remove(self, key: str, meta_only=False) -> Self:
+        if self.is_active() and key in self:
+            if self.is_read_only():
+                raise TypeError('Cache is in read only mode')
+            self._catalog.remove(key)
+            self._reader_writer.delete(key, meta_only=meta_only)
+        return self
+
+    def remove_meta(self, key: str) -> Self:
+        return self.remove(key, meta_only=True)
+
+    def keys(self) -> list[str]:
+        if not self.is_active():
+            return []
+        else:
+            return list(self._catalog.keys())
+
+    def _repr_params(self) -> list[str]:
+        params = [f'{len(self.keys())} items']
+        if not self.is_active():
+            params.append('NOT ACTIVE')
+        if self.is_read_only():
+            params.append('READ ONLY')
+        return params
+
+    def __repr__(self):
+        return f"{type(self).__name__}({', '.join(self._repr_params())})"
+
+
+def check_meta_mismatches(caches: Iterable[CacheCore], key=None) -> dict[str, tuple[CachedItemMeta, ...]]:
+    if key is not None and not isinstance(key, str):
+        raise TypeError('key is not a string')
+    for cache in caches:
+        if not isinstance(cache, CacheCore):
+            raise TypeError('caches must be an iterable containing Cache instances')
+    mismatches = {}
+    if key is None:
+        keys = set(k for cache in caches for k in cache.keys())
+    else:
+        keys = [key]
+    for k in keys:
+        metas = [cache.get_meta(k) for cache in caches if k in cache]
+        if len(metas) > 1:
+            first_meta = metas[0]
+            if any(meta != first_meta for meta in metas[1:]):
+                mismatches[k] = tuple(metas)
+    return mismatches
+
+
+class Serializer(metaclass=ABCMeta):
+    default_ext = None
+
+    def __init__(self, ext=None):
+        if ext is None:
+            self.ext = self.default_ext
+        else:
+            self.ext = ext
+
+    @abstractmethod
+    def read(self, path):
+        pass
+
+    @abstractmethod
+    def write(self, path, obj):
+        pass

yes3/caching/local_cache.py

@@ -0,0 +1,236 @@
+import json
+import os
+import pickle
+from datetime import datetime, UTC
+from functools import partial
+from glob import glob
+from pathlib import Path
+from typing import Optional, Self
+
+from yes3.caching.base import Cache, CacheDictCatalog, CachedItemMeta, Serializer, CacheReaderWriter
+
+
+class PickleSerializer(Serializer):
+    default_ext = 'pkl'
+
+    def read(self, path):
+        with open(path, 'rb') as f:
+            return pickle.load(f)
+
+    def write(self, path, obj):
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+        with open(path, 'wb') as f:
+            pickle.dump(obj, f)
+
+
+class JsonSerializer(Serializer):
+    default_ext = 'json'
+
+    def read(self, path) -> dict:
+        with open(path, 'r') as f:
+            return json.load(f)
+
+    def write(self, path, obj: dict):
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+        with open(path, 'w') as f:
+            json.dump(obj, f)
+
+
+class JsonMetaSerializer(JsonSerializer):
+    default_ext = 'meta'
+
+    def read(self, path) -> CachedItemMeta:
+        meta_dict = super().read(path)
+        return CachedItemMeta(**meta_dict)
+
+    def write(self, path, meta: CachedItemMeta):
+        super().write(path, meta.to_dict())
+
+
+def _get_serializer(serializer: str | Serializer, ext=None) -> Serializer:
+    if isinstance(serializer, type):
+        serializer = serializer(ext)
+
+    if isinstance(serializer, str):
+        if serializer.lstrip('.').lower() in {'pkl', 'pickle'}:
+            return PickleSerializer(ext)
+        elif serializer.lstrip('.').lower() == 'json':
+            return JsonSerializer(ext)
+        else:
+            raise NotImplementedError(f"Serializer not implemented for file type '{serializer}'")
+    elif isinstance(serializer, Serializer):
+        if ext is not None:
+            serializer.ext = ext
+        return serializer
+    else:
+        raise TypeError(
+            f'file_type must be a string or a Serializer subclass, but got type {type(serializer).__name__}')
+
+
+def _with_ext(path, ext: Optional[str]):
+    if ext is None:
+        return path
+    if not ext.startswith('.'):
+        ext = f'.{ext}'
+    path_str = str(path)
+    if path_str.endswith(ext):
+        return path
+    else:
+        try:
+            return type(path)(path_str + ext)
+        except (ValueError, TypeError):
+            return Path(path_str + ext)
+
+
+class LocalReaderWriter(CacheReaderWriter):
+    def __init__(
+            self, path: str | Path,
+            object_serializer: str | Serializer = PickleSerializer(),
+            meta_serializer: str | Serializer = JsonMetaSerializer(),
+    ):
+        self.path = Path(path)
+        self.obj_serializer = _get_serializer(object_serializer)
+        self.meta_serializer = _get_serializer(meta_serializer)
+
+    def clone(self, path: str | Path) -> Self:
+        return type(self)(path, object_serializer=self.obj_serializer, meta_serializer=self.meta_serializer)
+
+    def key2path(self, key: str, meta=False) -> Path:
+        if meta:
+            return self.path / _with_ext(key, self.meta_serializer.ext)
+        else:
+            return self.path / _with_ext(key, self.obj_serializer.ext)
+
+    def path2key(self, path: str | Path) -> str:
+        path = Path(path)
+        rel_path = path.relative_to(self.path)
+        return rel_path.stem
+
+    def read(self, key: str):
+        path = self.key2path(key)
+        print(f"Reading cached item '{key}' at {path}")
+        return self.obj_serializer.read(path)
+
+    def _build_meta(self, path, key=None) -> CachedItemMeta:
+        if key is None:
+            key = self.path2key(path)
+        file_stat = os.stat(path)
+        rel_path = path.relative_to(self.path)
+        return CachedItemMeta(
+            key=key,
+            path=str(rel_path),
+            size=file_stat.st_size,
+            timestamp=datetime.fromtimestamp(file_stat.st_mtime, UTC),
+        )
+
+    def get_meta(self, key: str, rebuild=False) -> CachedItemMeta:
+        if rebuild:
+            obj_path = self.key2path(key)
+            meta_path = self.key2path(key, meta=True)
+            meta = self._build_meta(path=obj_path, key=key)
+            self.meta_serializer.write(meta_path, meta)
+        else:
+            meta_path = self.key2path(key, meta=True)
+            meta = self.meta_serializer.read(meta_path)
+        return meta
+
+    def write(self, key: str, obj, meta: Optional[CachedItemMeta] = None) -> CachedItemMeta:
+        obj_path = self.key2path(key)
+        print(f"Caching item '{key}' at {obj_path}")
+        self.obj_serializer.write(obj_path, obj)
+
+        meta_path = self.key2path(key, meta=True)
+        if meta is None:
+            meta = self._build_meta(path=obj_path, key=key)
+        self.meta_serializer.write(meta_path, meta)
+        return meta
+
+    def delete(self, key: str, meta_only=False):
+        path = self.key2path(key)
+        meta_path = self.key2path(key, meta=True)
+        if meta_only:
+            print(f"Deleting cached item '{key}' metadata at {meta_path}")
+        else:
+            print(f"Deleting cached item '{key}' at {path}")
+            os.remove(path)
+        os.remove(meta_path)
+
+
+class LocalDiskCache(Cache):
+    @staticmethod
+    def _build_catalog_dict(reader_writer: LocalReaderWriter, rebuild_missing_meta=False) -> dict:
+        catalog_dict = {}
+        if os.path.exists(reader_writer.path):
+            data_ext = reader_writer.obj_serializer.ext.lstrip('.')
+            meta_ext = reader_writer.meta_serializer.ext.lstrip('.')
+            data_files = glob(str(reader_writer.path / f'*.{data_ext}'))
+            meta_files = glob(str(reader_writer.path / f'*.{meta_ext}'))
+            data_map = {Path(p).stem: p for p in data_files}
+            meta_map = {Path(p).stem: p for p in meta_files}
+            if data_map.keys() != meta_map.keys():
+                if rebuild_missing_meta:
+                    print(f'WARNING: data and metadata files are not aligned for cache at {reader_writer.path}, '
+                          'rebuilding missing metadata files')
+                else:
+                    raise RuntimeError(f'data and metadata files are not aligned for cache at {reader_writer.path}')
+            for key in data_map.keys():
+                catalog_dict[key] = reader_writer.get_meta(key, rebuild=(key not in meta_map and rebuild_missing_meta))
+        if len(catalog_dict.keys()) > 0:
+            print(f'{len(catalog_dict.keys())} cached items discovered at {reader_writer.path}')
+        return catalog_dict
+
+    @classmethod
+    def create(
+            cls,
+            path: str | Path,
+            obj_serializer: str | Serializer = PickleSerializer(),
+            meta_serializer: str | Serializer = JsonMetaSerializer(),
+            reader_writer: Optional[CacheReaderWriter] = None,
+            rebuild_missing_meta=False,
+            **kwargs,
+    ):
+        if reader_writer is None:
+            reader_writer = LocalReaderWriter(path, obj_serializer, meta_serializer)
+        elif not isinstance(reader_writer, LocalReaderWriter):
+            raise TypeError(f'`reader_writer` must be a {LocalReaderWriter.__name__} instance')
+        elif reader_writer.path != path:
+            reader_writer = reader_writer.clone(path)
+        catalog_builder = partial(cls._build_catalog_dict, reader_writer=reader_writer,
+                                  rebuild_missing_meta=rebuild_missing_meta)
+        catalog = CacheDictCatalog(catalog_builder=catalog_builder)
+        return cls(catalog, reader_writer, **kwargs)
+
+    @property
+    def path(self) -> Path:
+        return self._reader_writer.path
+
+    def subcache(self, rel_path: str) -> Self:
+        path = self.path / rel_path
+        kwargs = dict(active=self.is_active(), read_only=self.is_read_only())
+        return type(self).create(path, reader_writer=self._reader_writer, **kwargs)
+
+    def clear(self, force=False) -> Self:
+        if self.is_active() and len(self.keys()) > 0:
+            if not force:
+                raise RuntimeError(f'Clearing this {type(self).__name__} ({self.path}) requires specifying force=True')
+            print(f'Deleting {len(self.keys())} item(s) from cache at {self.path}')
+            for key in self.keys():
+                self.remove(key)
+            new_cache = type(self).create(self.path, reader_writer=self._reader_writer)
+            self.__init__(new_cache._catalog, new_cache._reader_writer, active=self._active, read_only=self._read_only)
+        return self
+
+    def clear_meta(self, force=False) -> Self:
+        if self.is_active() and len(self.keys()) > 0:
+            if not force:
+                raise RuntimeError(f'Clearing this {type(self).__name__} metadata ({self.path}) requires specifying '
+                                   'force=True')
+            print(f'Deleting {len(self.keys())} item(s) from cache at {self.path}')
+            for key in self.keys():
+                self.remove(key, meta_only=True)
+        return self
+
+    def _repr_params(self) -> list[str]:
+        params = super()._repr_params()
+        params.insert(0, str(self.path))
+        return params

yes3/caching/memory_cache.py

@@ -0,0 +1,74 @@
+from datetime import datetime, UTC
+from typing import Any, Optional, Self
+
+from yes3.caching.base import CacheCore, CachedItemMeta, raise_not_found, _NotSpecified
+
+
+class MemoryCache(CacheCore):
+    def __init__(self, active=True, read_only=False):
+        super().__init__(active=active, read_only=read_only)
+        self._data: dict[str, Any] = {}
+        self._meta: dict[str, CachedItemMeta] = {}
+
+    def __contains__(self, key: str):
+        if not self.is_active():
+            return False
+        found = key in self._data
+        if found and key not in self._meta:
+            raise RuntimeError(f"data exists, but no metadata found, for key '{key}' in {type(self).__name__}")
+        return found
+
+    def get(self, key: str, default=_NotSpecified):
+        if not self.is_active() or key not in self:
+            if default is _NotSpecified:
+                raise_not_found(key)
+            else:
+                return default
+        return self._data[key]
+
+    def get_meta(self, key: str) -> CachedItemMeta:
+        if not self.is_active() or key not in self:
+            raise_not_found(key)
+        return self._meta[key]
+
+    def put(self, key: str, obj, *, update=False, meta: Optional[CachedItemMeta] = None) -> Self:
+        if self.is_read_only():
+            raise TypeError('Cache is in read only mode')
+        if self.is_active():
+            if key in self and not update:
+                raise ValueError(f"key '{key}' already exists in cache; use 'update' to overwrite")
+            if meta is None:
+                meta = CachedItemMeta(key=key, timestamp=datetime.now(UTC), path=None, size=None)
+            self._meta[key] = meta
+            self._data[key] = obj
+        else:
+            print(f"WARNING: {type(self).__name__} is not active")
+        return self
+
+    def remove(self, key: str) -> Self:
+        if self.is_active():
+            if key in self:
+                if self.is_read_only():
+                    raise TypeError('Cache is in read only mode')
+                self._data.pop(key)
+                self._meta.pop(key)
+        else:
+            print(f"WARNING: {type(self).__name__} is not active")
+        return self
+
+    def keys(self) -> list[Any]:
+        if not self.is_active():
+            return []
+        else:
+            return list(self._data.keys())
+
+    def clear(self, force=False) -> Self:
+        if self.is_active():
+            if len(self.keys()) > 0:
+                if not force:
+                    raise RuntimeError(f'Clearing this {type(self).__name__} requires specifying force=True')
+                self._data: dict[str, Any] = {}
+                self._meta: dict[str, CachedItemMeta] = {}
+        else:
+            print(f"WARNING: {type(self).__name__} is not active")
+        return self