Added key_attribute for decorators

aiocache/utils.py

@@ -8,7 +8,7 @@
 logger = logging.getLogger(__name__)
 
 
-def cached(*args, ttl=0, backend=None, serializer=None, **kwargs):
+def cached(*args, ttl=0, key_attribute=None, backend=None, serializer=None, **kwargs):
     """
     Caches the functions return value into a key generated with module_name, function_name and args.
 
@@ -17,6 +17,8 @@ def cached(*args, ttl=0, backend=None, serializer=None, **kwargs):
     to the backend class when instantiating.
 
     :param ttl: int seconds to store the function call. Default is 0
+    :param key_attribute: keyword attribute from the function to use as a key. If not passed,
+        it will use module_name + function_name + args + kwargs
     :param backend: backend class to use when calling the ``set``/``get`` operations. Default is
         :class:`aiocache.backends.SimpleMemoryCache`
     :param serializer: serializer instance to use when calling the ``serialize``/``deserialize``.
@@ -26,7 +28,8 @@ def cached(*args, ttl=0, backend=None, serializer=None, **kwargs):
 
     def cached_decorator(fn):
         async def wrapper(*args, **kwargs):
-            key = (fn.__module__ or "stub") + fn.__name__ + str(args) + str(kwargs)
+            key = kwargs.get(
+                key_attribute, (fn.__module__ or 'stub') + fn.__name__ + str(args) + str(kwargs))
             if await cache.exists(key):
                 return await cache.get(key)
             else:
@@ -37,7 +40,7 @@ async def wrapper(*args, **kwargs):
     return cached_decorator
 
 
-def multi_cached(*args, backend=None, serializer=None, **kwargs):
+def multi_cached(*args, key_attribute=None, backend=None, serializer=None, **kwargs):
     """
     Only supports functions that return dict-like structures. This decorator caches each key/value
     of the dict-like object returned by the function.
@@ -47,6 +50,9 @@ def multi_cached(*args, backend=None, serializer=None, **kwargs):
             returned in the response. If its not the case, the call to the function will always
             be done (although the returned values will all be cached).
 
+    :param key_attribute: keyword attribute from the function containing an iterable to use
+        as a keys. If not passed, it will try with 'keys' attribute and if none of them exists,
+        it won't try to reuse keys from the cache
     :param backend: backend class to use when calling the ``set``/``get`` operations. Default is
         :class:`aiocache.backends.SimpleMemoryCache`
     :param serializer: serializer instance to use when calling the ``serialize``/``deserialize``.
@@ -58,7 +64,7 @@ def multi_cached_decorator(fn):
         async def wrapper(*args, **kwargs):
             partial_dict = {}
             missing_keys = []
-            keys = (kwargs.pop("keys", []))
+            keys = kwargs.get(key_attribute, [])
             if keys:
                 values = await cache.multi_get(keys)
                 for key, value in zip(keys, values):

tests/integration/test_utils.py → tests/ut/test_utils.py

@@ -1,9 +1,11 @@
 import sys
 import pytest
 import aiocache
-import asyncio
+import random
 
-from aiocache import RedisCache, SimpleMemoryCache, cached, multi_cached
+from unittest import mock
+
+from aiocache import RedisCache, SimpleMemoryCache, cached, multi_cached, config_default_cache
 from aiocache.utils import get_default_cache
 from aiocache.serializers import PickleSerializer, DefaultSerializer
 
@@ -14,27 +16,61 @@ async def return_dict(*args, keys=None):
         ret[key] = value
     return ret
 
+async def stub(*args, **kwargs):
+    return random.randint(1, 50)
+
+
+@pytest.fixture
+def mock_cache(mocker):
+    config_default_cache()
+    cache = get_default_cache()
+    mocker.spy(cache, 'multi_set')
+    mocker.spy(cache, 'multi_get')
+    mocker.spy(cache, 'get')
+    mocker.spy(cache, 'exists')
+    mocker.spy(cache, 'set')
+    return cache
+
 
 class TestCachedDecorator:
 
     @pytest.mark.asyncio
-    async def test_cached_no_args(self, mocker):
-        mocker.spy(asyncio, 'sleep')
+    async def test_cached_ttl(self, mocker, mock_cache):
+        module = sys.modules[globals()['__name__']]
+        mocker.spy(module, 'stub')
         cached_decorator = cached(ttl=10)
 
-        resp1 = await cached_decorator(asyncio.sleep)(1)
-        resp2 = await cached_decorator(asyncio.sleep)(1)
+        resp1 = await cached_decorator(stub)(1)
+        resp2 = await cached_decorator(stub)(1)
 
-        assert asyncio.sleep.call_count == 1
+        assert stub.call_count == 1
         assert resp1 is resp2
 
+        mock_cache.get.assert_called_with('stubstub(1,){}')
+        assert mock_cache.get.call_count == 1
+        assert mock_cache.exists.call_count == 2
+        mock_cache.set.assert_called_with('stubstub(1,){}', mock.ANY, ttl=10)
+        assert mock_cache.set.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_cached_key_attribute(self, mocker, mock_cache):
+        module = sys.modules[globals()['__name__']]
+        mocker.spy(module, 'stub')
+        cached_decorator = cached(key_attribute="key")
+
+        await cached_decorator(stub)(key='key')
+        await cached_decorator(stub)(key='key')
+
+        mock_cache.get.assert_called_with('key')
+        mock_cache.set.assert_called_with('key', mock.ANY, ttl=0)
+
 
 class TestMultiCachedDecorator:
     @pytest.mark.asyncio
-    async def test_multi_cached_no_args(self, mocker):
+    async def test_multi_cached(self, mocker):
         module = sys.modules[globals()['__name__']]
         mocker.spy(module, 'return_dict')
-        cached_decorator = multi_cached()
+        cached_decorator = multi_cached(key_attribute='keys')
 
         default_keys = {'a', 'd', 'z', 'y'}
         resp_default = await cached_decorator(return_dict)()
@@ -51,6 +87,26 @@ async def test_multi_cached_no_args(self, mocker):
         return_dict.assert_called_with(keys=list(keys2 - keys1 - default_keys))
         assert keys2 == set(resp2.keys())
 
+    @pytest.mark.asyncio
+    async def test_multi_cached_key_attribute(self, mocker, mock_cache):
+        module = sys.modules[globals()['__name__']]
+        mocker.spy(module, 'return_dict')
+        cached_decorator = multi_cached(key_attribute='keys')
+        keys1 = {'a', 'b'}
+
+        await cached_decorator(return_dict)(keys=keys1)
+        mock_cache.multi_get.assert_called_once_with(keys1)
+        mock_cache.multi_set.call_count = 1
+
+    @pytest.mark.asyncio
+    async def test_multi_cached_no_key_attribute(self, mocker, mock_cache):
+        module = sys.modules[globals()['__name__']]
+        mocker.spy(module, 'return_dict')
+        cached_decorator = multi_cached()
+
+        await cached_decorator(return_dict)()
+        mock_cache.multi_set.call_count = 1
+
 
 class TestDefaultCache: