"""Cache management""" from abc import ABC, abstractmethod from pathlib import Path from pickle import dumps, loads from typing import Awaitable, Callable, Generic, TypeVar, Union, cast from xdg import xdg_cache_home ResourceType = TypeVar("ResourceType") CacheFallback = Union[ResourceType, Callable[[str], Awaitable[ResourceType]]] class Cache(Generic[ResourceType], ABC): """Base class for caches""" @abstractmethod async def get( self, key: str, fallback: CacheFallback[ResourceType] ) -> ResourceType: """Get an item in the cache, call fallback if it's not present""" @abstractmethod def set(self, key: str, resource: ResourceType) -> None: """Set a resource in the cache""" @staticmethod async def _get_fallback_value( key: str, fallback: CacheFallback[ResourceType] ) -> ResourceType: if callable(fallback): result = await fallback(key) else: result = fallback return result class NullCache(Cache[ResourceType]): """Disabled cache""" async def get( self, key: str, fallback: CacheFallback[ResourceType] ) -> ResourceType: return await self._get_fallback_value(key, fallback) def set(self, key: str, resource: ResourceType) -> None: pass class FileCache(Cache[ResourceType]): """Cache on the local filesystem""" def __init__(self, name: str) -> None: self._name = name async def get( self, key: str, fallback: CacheFallback[ResourceType] ) -> ResourceType: """Get an item in the cache, call fallback if it's not present""" cache_file_path = self._get_cache_file_path(key) if not cache_file_path.is_file(): resource = await self._get_fallback_value(key, fallback) self.set(key, resource) else: with open(cache_file_path, "rb") as cache_file: resource_data = cache_file.read() resource = self._deserialize(resource_data) return resource def set(self, key: str, resource: ResourceType) -> None: """Set a resource in the cache""" cache_file_path = self._get_cache_file_path(key) with open(cache_file_path, "wb") as cache_file: resource_data = self._serialize(resource) cache_file.write(resource_data) def _get_cache_file_path(self, key: str) -> Path: key_slug = _get_key_slug(key) cache_directory = xdg_cache_home() / "frontools" / self._name file_path = cache_directory.joinpath(*key_slug.split("&")) file_directory = file_path.parent if not file_directory.is_dir(): file_directory.mkdir(parents=True) return file_path @abstractmethod def _serialize(self, resource: ResourceType) -> bytes: pass @abstractmethod def _deserialize(self, data: bytes) -> ResourceType: pass class DataCache(FileCache[bytes]): """Cache of byte array""" def _serialize(self, resource: bytes) -> bytes: return resource def _deserialize(self, data: bytes) -> bytes: return data class ObjectCache(FileCache[ResourceType]): """Cache pickling objects""" def _serialize(self, resource: ResourceType) -> bytes: return dumps(resource) def _deserialize(self, data: bytes) -> ResourceType: return cast(ResourceType, loads(data)) def _get_key_slug(url: str) -> str: """Return an unique slug usable as a path name for a given url.""" return url.replace("_", "___").replace("/", "__").replace(":", "_")