misc-csechet/frontools/cache.py

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