multi_cached can use arg for key_from_attr

aiocache/decorators.py

@@ -174,7 +174,7 @@ def _get_cache(
 def _get_args_dict(func, args, kwargs):
     defaults = {
         arg_name: arg.default for arg_name, arg in inspect.signature(func).parameters.items()
-        if arg.default is not inspect._empty
+        if arg.default is not inspect._empty   # TODO: bug prone..
     }
     args_names = func.__code__.co_varnames[:func.__code__.co_argcount]
     return {**defaults, **dict(zip(args_names, args)), **kwargs}
@@ -239,7 +239,7 @@ async def wrapper(*args, **kwargs):
     async def decorator(self, f, *args, **kwargs):
         missing_keys = []
         partial = {}
-        keys = self.get_cache_keys(f, args, kwargs)
+        keys, new_args = self.get_cache_keys(f, args, kwargs)
 
         values = await self.get_from_cache(*keys)
         for key, value in zip(keys, values):
@@ -251,7 +251,7 @@ async def decorator(self, f, *args, **kwargs):
         if values and None not in values:
             return partial
 
-        result = await f(*args, **kwargs)
+        result = await f(*new_args, **kwargs)
         result.update(partial)
         await self.set_in_cache(result, args, kwargs)
 
@@ -260,7 +260,14 @@ async def decorator(self, f, *args, **kwargs):
     def get_cache_keys(self, f, args, kwargs):
         args_dict = _get_args_dict(f, args, kwargs)
         keys = args_dict[self.keys_from_attr] or []
-        return [self.key_builder(key, *args, **kwargs) for key in keys]
+        keys = [self.key_builder(key, *args, **kwargs) for key in keys]
+
+        args_names = f.__code__.co_varnames[:f.__code__.co_argcount]
+        new_args = list(args)
+        if self.keys_from_attr in args_names and self.keys_from_attr not in kwargs:
+            new_args[args_names.index(self.keys_from_attr)] = keys
+
+        return keys, tuple(new_args)
 
     async def get_from_cache(self, *keys):
         if not keys:

tests/acceptance/test_decorators.py

@@ -120,3 +120,40 @@ async def fn(self, keys, market='ES'):
         await fn('self', keys=['a', 'b'])
         assert await cache.exists('a_ES') is True
         assert await cache.exists('b_ES') is True
+
+    @pytest.mark.asyncio
+    async def test_fn_with_args(self, cache):
+
+        @multi_cached('keys')
+        async def fn(keys, *args):
+            assert len(args) == 1
+            return {'a': 1}
+
+        await fn(['a'], 1)
+        assert await cache.exists('a') is True
+
+    @pytest.mark.asyncio
+    async def test_keys_without_kwarg(self, cache):
+
+        @multi_cached('keys')
+        async def fn(keys):
+            return {'a': 1}
+
+        await fn(['a'])
+        assert await cache.exists('a') is True
+
+    @pytest.mark.asyncio
+    async def test_double_decorator(self, cache):
+
+        def dummy_d(fn):
+            async def wrapper(*args, **kwargs):
+                await fn(*args, **kwargs)
+            return wrapper
+
+        @dummy_d
+        @multi_cached('keys')
+        async def fn(keys):
+            return {'a': 1}
+
+        await fn(['a'])
+        assert await cache.exists('a') is True

tests/ut/test_decorators.py

@@ -9,6 +9,7 @@
 
 from aiocache.base import BaseCache
 from aiocache import cached, cached_stampede, multi_cached, SimpleMemoryCache
+from aiocache.decorators import _get_args_dict
 from aiocache.serializers import JsonSerializer
 
 
@@ -322,22 +323,28 @@ def test_alias_takes_precedence(self, mock_cache):
             assert mc.cache is mock_cache
 
     def test_get_cache_keys(self, decorator):
-        assert decorator.get_cache_keys(stub_dict, (), {'keys': ['a', 'b']}) == ['a', 'b']
+        assert decorator.get_cache_keys(
+            stub_dict, (), {'keys': ['a', 'b']}) == (['a', 'b'], ())
 
     def test_get_cache_keys_empty_list(self, decorator):
-        assert decorator.get_cache_keys(stub_dict, (), {'keys': []}) == []
+        assert decorator.get_cache_keys(stub_dict, (), {'keys': []}) == ([], ())
 
     def test_get_cache_keys_missing_kwarg(self, decorator):
         with pytest.raises(KeyError):
             assert decorator.get_cache_keys(stub_dict, (), {})
 
+    def test_get_cache_keys_arg_key_from_attr(self, decorator):
+        def fake(keys, a=1, b=2):
+            pass
+        assert decorator.get_cache_keys(fake, (['a']), {}) == (['a'], (['a'],))
+
     def test_get_cache_keys_with_none(self, decorator):
-        assert decorator.get_cache_keys(stub_dict, (), {'keys': None}) == []
+        assert decorator.get_cache_keys(stub_dict, (), {'keys': None}) == ([], ())
 
     def test_get_cache_keys_with_key_builder(self, decorator):
         decorator.key_builder = lambda key, *args, **kwargs: kwargs['market'] + '_' + key.upper()
         assert decorator.get_cache_keys(
-            stub_dict, (), {'keys': ['a', 'b'], 'market': 'ES'}) == ['ES_A', 'ES_B']
+            stub_dict, (), {'keys': ['a', 'b'], 'market': 'ES'}) == (['ES_A', 'ES_B'], ())
 
     @pytest.mark.asyncio
     async def test_get_from_cache(self, decorator, decorator_call):
@@ -472,3 +479,11 @@ async def what(keys=None):
 
             assert get_c.call_count == 1
             assert cache.multi_get.call_count == 2
+
+
+def test_get_args_dict():
+    def fn(a, b, *args, keys=None, **kwargs):
+        pass
+
+    args_dict = _get_args_dict(fn, ('a', 'b', 'c', 'd'), {'what': 'what'})
+    assert args_dict == {'a': 'a', 'b': 'b', 'keys': None, 'what': 'what'}