cache.py

from collections import deque

import psutil

from settings import settings
from collections import AsyncIterable


class AsyncIteratorWrapper(AsyncIterable):
    def __init__(self, obj):
        self._it = iter(obj)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            value = next(self._it)
        except StopIteration:
            raise StopAsyncIteration
        return value


class Cache:
    def __init__(self, items_per_page=100000, max_pages=5):
        self.items_per_page = items_per_page
        self.max_pages = max_pages
        self._pages = deque()

    async def clean(self):
        while len(self._pages) > self.max_pages:
            self._pages.popleft()

        if (len(self._pages) > 1 and
                self._free_space() <= settings.MIN_FREE_MEMORY_SPACE_IN_BYTES):
            self._pages.popleft()

        while (len(self._pages) > 1 and
               self._free_space_percent() <=
               settings.MIN_FREE_MEMORY_SPACE_PERCENT):
            self._pages.popleft()

    def _free_space(self) -> float:
        mem = psutil.virtual_memory()
        return mem.available

    def _free_space_percent(self) -> float:
        mem = psutil.virtual_memory()
        return 1 - (mem.percent / 100)

    async def add(self, key: str, value: dict) -> dict:
        if (len(self._pages) == 0 or
                len(self._pages[-1]) >= self.items_per_page):
            self._new_page({key: value})
            await self.clean()
            return value
        self._pages[-1][key] = value
        return value

    def _new_page(self, page: dict):
        self._pages.append(page)

    async def get(self, key) -> dict:
        async for page in AsyncIteratorWrapper(self._pages):
            try:
                return page[key]
            except KeyError:
                pass
        raise KeyError()

    def __len__(self):
        return sum(len(page) for page in self._pages)

    def _number_of_pages(self):
        return len(self._pages)

    def reset(self):
        self._pages = deque()