PyNyaaTa/pynyaata2/cache/simple.py

from datetime import datetime, timedelta
from os import getenv
from typing import Dict, List, Optional, Tuple

from pynyaata2.types import Cache, RemoteFile


CACHE_TIMEOUT = int(getenv("CACHE_TIMEOUT", 60 * 60))
CACHE_DATA: Dict[str, Tuple[List[RemoteFile], datetime]] = {}


class SimpleCache(Cache):
    def get(self, key: str) -> Optional[List[RemoteFile]]:
        if key in CACHE_DATA:
            data, timeout = CACHE_DATA[key]

            if datetime.now() < timeout + timedelta(seconds=CACHE_TIMEOUT):
                return data
            else:
                CACHE_DATA.pop(key)

        return None

    def set(self, key: str, data: List[RemoteFile]):
        CACHE_DATA[key] = (data, datetime.now())