first commit

fastapi_cache/__init__.py

@@ -0,0 +1,18 @@
+class FastAPICache:
+    _backend = None
+    _prefix = None
+
+    @classmethod
+    def init(cls, backend, prefix: str = ""):
+        cls._backend = backend
+        cls._prefix = prefix
+
+    @classmethod
+    def get_backend(cls):
+        assert cls._backend, "You must call init first!"  # nosec: B101
+        return cls._backend
+
+    @classmethod
+    def get_prefix(cls):
+        assert cls._prefix, "You must call init first!"  # nosec: B101
+        return cls._prefix

fastapi_cache/backends/__init__.py

@@ -0,0 +1,16 @@
+import abc
+from typing import Tuple
+
+
+class Backend:
+    @abc.abstractmethod
+    async def get_with_ttl(self, key: str) -> Tuple[int, str]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get(self, key: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def set(self, key: str, value: str, expire: int = None):
+        raise NotImplementedError

fastapi_cache/backends/mencache.py

@@ -0,0 +1,19 @@
+from typing import Tuple
+
+from aiomcache import Client
+
+from fastapi_cache.backends import Backend
+
+
+class MemcacheBackend(Backend):
+    def __init__(self, mcache: Client):
+        self.mcache = mcache
+
+    async def get_with_ttl(self, key: str) -> Tuple[int, str]:
+        return 0, await self.mcache.get(key.encode())
+
+    async def get(self, key: str):
+        return await self.mcache.get(key, key.encode())
+
+    async def set(self, key: str, value: str, expire: int = None):
+        return await self.mcache.set(key.encode(), value.encode(), exptime=expire)

fastapi_cache/backends/redis.py

@@ -0,0 +1,22 @@
+from typing import Tuple
+
+from aioredis import Redis
+
+from fastapi_cache.backends import Backend
+
+
+class RedisBackend(Backend):
+    def __init__(self, redis: Redis):
+        self.redis = redis
+
+    async def get_with_ttl(self, key: str) -> Tuple[int, str]:
+        p = self.redis.pipeline()
+        p.ttl(key)
+        p.get(key)
+        return await p.execute()
+
+    async def get(self, key) -> str:
+        return await self.redis.get(key)
+
+    async def set(self, key: str, value: str, expire: int = None):
+        return await self.redis.set(key, value, expire=expire)

fastapi_cache/coder.py

@@ -0,0 +1,33 @@
+import json
+import pickle  # nosec:B403
+from typing import Any
+
+
+class Coder:
+    @classmethod
+    def encode(cls, value: Any):
+        raise NotImplementedError
+
+    @classmethod
+    def decode(cls, value: Any):
+        raise NotImplementedError
+
+
+class JsonCoder(Coder):
+    @classmethod
+    def encode(cls, value: Any):
+        return json.dumps(value)
+
+    @classmethod
+    def decode(cls, value: Any):
+        return json.loads(value)
+
+
+class PickleCoder(Coder):
+    @classmethod
+    def encode(cls, value: Any):
+        return pickle.dumps(value)
+
+    @classmethod
+    def decode(cls, value: Any):
+        return pickle.loads(value)  # nosec:B403

fastapi_cache/decorator.py

@@ -0,0 +1,99 @@
+from functools import wraps
+from typing import Callable, Optional, Type
+
+from starlette.requests import Request
+from starlette.responses import Response
+
+from fastapi_cache import FastAPICache
+from fastapi_cache.coder import Coder, JsonCoder
+
+
+def default_key_builder(
+    func,
+    namespace: Optional[str] = "",
+    request: Request = None,
+    response: Response = None,
+    *args,
+    **kwargs,
+):
+    prefix = FastAPICache.get_prefix()
+    cache_key = f"{prefix}:{namespace}:{func.__module__}:{func.__name__}:{args}:{kwargs}"
+    return cache_key
+
+
+def cache(
+    expire: int = None,
+    coder: Type[Coder] = JsonCoder,
+    key_builder: Callable = default_key_builder,
+    namespace: Optional[str] = "",
+):
+    """
+    cache all function
+    :param namespace:
+    :param expire:
+    :param coder:
+    :param key_builder:
+    :return:
+    """
+
+    def wrapper(func):
+        @wraps(func)
+        async def inner(*args, **kwargs):
+            backend = FastAPICache.get_backend()
+            cache_key = key_builder(func, namespace, *args, **kwargs)
+            ret = await backend.get(cache_key)
+            if ret is not None:
+                return coder.decode(ret)
+
+            ret = await func(*args, **kwargs)
+            await backend.set(cache_key, coder.encode(ret), expire)
+            return ret
+
+        return inner
+
+    return wrapper
+
+
+def cache_response(
+    expire: int = None,
+    coder: Type[Coder] = JsonCoder,
+    key_builder: Callable = default_key_builder,
+    namespace: Optional[str] = "",
+):
+    """
+    cache fastapi response
+    :param namespace:
+    :param expire:
+    :param coder:
+    :param key_builder:
+    :return:
+    """
+
+    def wrapper(func):
+        @wraps(func)
+        async def inner(request: Request, *args, **kwargs):
+            if request.method != "GET":
+                return await func(request, *args, **kwargs)
+
+            backend = FastAPICache.get_backend()
+            cache_key = key_builder(func, namespace, request, *args, **kwargs)
+            ttl, ret = await backend.get_with_ttl(cache_key)
+            if_none_match = request.headers.get("if-none-match")
+            if ret is not None:
+                response = kwargs.get("response")
+                if response:
+                    response.headers["Cache-Control"] = f"max-age={ttl}"
+                    etag = f"W/{hash(ret)}"
+                    if if_none_match == etag:
+                        response.status_code = 304
+                        return response
+                    response.headers["ETag"] = etag
+                return coder.decode(ret)
+
+            ret = await func(request, *args, **kwargs)
+            await backend.set(cache_key, coder.encode(ret), expire)
+            return ret
+
+        return inner
+
+    return wrapper