misc-csechet/frontools/cache.py

122 lines
3.5 KiB
Python

"""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(":", "_")