blob_lru.py

# Copyright (c) 2020, Andrew Zhou
# All rights reserved.

from __future__ import annotations

from aiosqlite import Connection as AConnection

from collections import OrderedDict
from inspect import Signature
from typing import Any, Callable, Generic, Optional, TypeVar
from typing import overload


KeyType = TypeVar('KeyType', int, str)
ReturnType = TypeVar('ReturnType')


@overload
def blob_lru(user_function: Callable) -> AbstractBlobCache:
    ...


@overload
def blob_lru(max_size: int = 128) -> Callable[[Callable], AbstractBlobCache]:
    ...


def blob_lru(func_or_size):
    if callable(func_or_size):
        return AbstractBlobCache(func_or_size)

    def deco(func: Callable):
        return AbstractBlobCache(func, max_size=func_or_size)

    return deco


class AbstractBlobCache(Generic[ReturnType, KeyType]):
    def __init__(self, func: Callable[..., ReturnType], cache_size: int = 128):
        self.cache_size = cache_size
        self.func = func
        self.signature = Signature.from_callable(func)
        self.cache: OrderedDict[KeyType, ReturnType] = OrderedDict()

    def _store(self, key: KeyType, val: ReturnType) -> None:
        self.cache[key] = val
        if len(self.cache) > self.cache_size:
            self.cache.popitem(last=False)

    def _fetch(self, key: KeyType) -> Optional[ReturnType]:
        if key not in self.cache:
            return None
        self.cache.move_to_end(key)
        return self.cache[key]

    def _key(self, arguments: OrderedDict) -> KeyType:
        raise NotImplementedError

    def __call__(self, *args: Any, **kwargs: Any) -> ReturnType:
        ba = self.signature.bind(*args, **kwargs)
        ba.apply_defaults()
        key = self._key(ba.arguments)
        val = self._fetch(key)
        if val is None:
            val = self.func(*ba.args, **ba.kwargs)
            self._store(key, val)
        return val


class AiosqliteBlobCache(AbstractBlobCache):
    def __init__(
        self,
        func: Callable[..., ReturnType],
        connection: AConnection,
        cache_size: int = 128,
        cold_cache_size: int = 512,
    ):
        # TODO: Finish this
        super().__init__(func, cache_size)
        self.connection = connection
        self.cold_cache_size = cold_cache_size