143 lines
4.3 KiB
Python
143 lines
4.3 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 click import echo
|
|
from shutil import rmtree
|
|
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"""
|
|
|
|
cache_base = xdg_cache_home() / "frontools"
|
|
|
|
@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
|
|
def prune(cache_names: list[str]) -> None:
|
|
"""Remove caches from filesystem.
|
|
|
|
If empty list is provided, all caches will be cleaned
|
|
"""
|
|
if not len(cache_names):
|
|
cache_names = [it.name for it in Cache.cache_base.iterdir() if it.is_dir()]
|
|
for cache_name in cache_names:
|
|
cache_path: Path = Cache.cache_base / cache_name
|
|
if not cache_path.is_dir():
|
|
echo(f'{cache_path} isn\'t a chache directory', err=True)
|
|
continue
|
|
echo(f'Removing {cache_path}')
|
|
rmtree(cache_path)
|
|
|
|
@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 = self.cache_base / self._name
|
|
file_path = cache_directory.joinpath(*key_slug.split("&"))
|
|
file_path = file_path.parent / (file_path.name + '_')
|
|
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(":", "_")
|