Add function as parameter for key_builder (#417)

aiocache/decorators.py

@@ -30,7 +30,7 @@ class cached:
         key_builder param. If key and key_builder are not passed, it will use module_name
         + function_name + args + kwargs
     :param key_builder: Callable that allows to build the function dynamically. It receives
-        same args and kwargs as the called function.
+        the function plus same args and kwargs passed to the function.
     :param cache: cache class to use when calling the ``set``/``get`` operations.
         Default is ``aiocache.SimpleMemoryCache``.
     :param serializer: serializer instance to use when calling the ``dumps``/``loads``.
@@ -106,7 +106,7 @@ def get_cache_key(self, f, args, kwargs):
         if self.key:
             return self.key
         if self.key_builder:
-            return self.key_builder(*args, **kwargs)
+            return self.key_builder(f, *args, **kwargs)
 
         return self._key_from_args(f, args, kwargs)
 
@@ -224,7 +224,7 @@ class multi_cached:
     :param keys_from_attr: arg or kwarg name from the function containing an iterable to use
         as keys to index in the cache.
     :param key_builder: Callable that allows to change the format of the keys before storing.
-        Receives the key and same args and kwargs as the called function.
+        Receives the key the function and same args and kwargs as the called function.
     :param ttl: int seconds to store the keys. Default is 0 which means no expiration.
     :param cache: cache class to use when calling the ``multi_set``/``multi_get`` operations.
         Default is ``aiocache.SimpleMemoryCache``.
@@ -249,7 +249,7 @@ def __init__(
         **kwargs
     ):
         self.keys_from_attr = keys_from_attr
-        self.key_builder = key_builder or (lambda key, *args, **kwargs: key)
+        self.key_builder = key_builder or (lambda key, f, *args, **kwargs: key)
         self.ttl = ttl
         self.alias = alias
         self.cache = None
@@ -303,14 +303,14 @@ async def decorator(self, f, *args, cache_read=True, cache_write=True, **kwargs)
         result.update(partial)
 
         if cache_write:
-            await self.set_in_cache(result, args, kwargs)
+            await self.set_in_cache(result, f, args, kwargs)
 
         return result
 
     def get_cache_keys(self, f, args, kwargs):
         args_dict = _get_args_dict(f, args, kwargs)
         keys = args_dict[self.keys_from_attr] or []
-        keys = [self.key_builder(key, *args, **kwargs) for key in keys]
+        keys = [self.key_builder(key, f, *args, **kwargs) for key in keys]
 
         args_names = f.__code__.co_varnames[: f.__code__.co_argcount]
         new_args = list(args)
@@ -331,10 +331,10 @@ async def get_from_cache(self, *keys):
             logger.exception("Couldn't retrieve %s, unexpected error", keys)
             return [None] * len(keys)
 
-    async def set_in_cache(self, result, fn_args, fn_kwargs):
+    async def set_in_cache(self, result, fn, fn_args, fn_kwargs):
         try:
             await self.cache.multi_set(
-                [(self.key_builder(k, *fn_args, **fn_kwargs), v) for k, v in result.items()],
+                [(self.key_builder(k, fn, *fn_args, **fn_kwargs), v) for k, v in result.items()],
                 ttl=self.ttl,
             )
         except Exception:

tests/acceptance/conftest.py

@@ -61,4 +61,4 @@ def memcached_cache(event_loop):
 
 @pytest.fixture(params=["redis_cache", "memory_cache", "memcached_cache"])
 def cache(request):
-    return request.getfuncargvalue(request.param)
+    return request.getfixturevalue(request.param)

tests/acceptance/test_decorators.py

@@ -41,15 +41,15 @@ async def fn():
 
     @pytest.mark.asyncio
     async def test_cached_key_builder(self, cache):
-        def build_key(self, a, b):
-            return "{}_{}_{}".format(self, a, b)
+        def build_key(f, self, a, b):
+            return "{}_{}_{}_{}".format(self, f.__name__, a, b)
 
         @cached(key_builder=build_key)
         async def fn(self, a, b=2):
             return "1"
 
         await fn("self", 1, 3)
-        assert await cache.exists(build_key("self", 1, 3)) is True
+        assert await cache.exists(build_key(fn, "self", 1, 3)) is True
 
 
 class TestCachedStampede:
@@ -121,16 +121,16 @@ async def fn(keys):
 
     @pytest.mark.asyncio
     async def test_multi_cached_key_builder(self, cache):
-        def build_key(key, self, keys, market="ES"):
-            return "{}_{}".format(key, market)
+        def build_key(key, f, self, keys, market="ES"):
+            return "{}_{}_{}".format(f.__name__, key, market)
 
         @multi_cached(keys_from_attr="keys", key_builder=build_key)
         async def fn(self, keys, market="ES"):
             return {pytest.KEY: 1, pytest.KEY_1: 2}
 
         await fn("self", keys=[pytest.KEY, pytest.KEY_1])
-        assert await cache.exists(pytest.KEY + "_ES") is True
-        assert await cache.exists(pytest.KEY_1 + "_ES") is True
+        assert await cache.exists("fn_" + pytest.KEY + "_ES") is True
+        assert await cache.exists("fn_" + pytest.KEY_1 + "_ES") is True
 
     @pytest.mark.asyncio
     async def test_fn_with_args(self, cache):

tests/ut/test_decorators.py

@@ -362,7 +362,7 @@ def test_init(self):
         )
 
         assert mc.ttl == 1
-        assert mc.key_builder("key") == "key"
+        assert mc.key_builder("key", lambda x: x) == "key"
         assert mc.keys_from_attr == "keys"
         assert mc.cache is None
         assert mc._cache == SimpleMemoryCache
@@ -468,7 +468,7 @@ async def test_calls_fn_multi_set_when_multi_get_none(self, mocker, decorator, d
         ret = await decorator_call(1, keys=["a", "b"], value="value")
 
         decorator.get_from_cache.assert_called_once_with("a", "b")
-        decorator.set_in_cache.assert_called_with(ret, ANY, ANY)
+        decorator.set_in_cache.assert_called_with(ret, stub_dict, ANY, ANY)
         stub_dict.assert_called_once_with(1, keys=["a", "b"], value="value")
 
     @pytest.mark.asyncio
@@ -478,7 +478,7 @@ async def test_calls_fn_with_only_missing_keys(self, mocker, decorator, decorato
 
         assert await decorator_call(1, keys=["a", "b"], value="value") == {"a": ANY, "b": ANY}
 
-        decorator.set_in_cache.assert_called_once_with({"a": ANY, "b": ANY}, ANY, ANY)
+        decorator.set_in_cache.assert_called_once_with({"a": ANY, "b": ANY}, stub_dict, ANY, ANY)
         stub_dict.assert_called_once_with(1, keys=["b"], value="value")
 
     @pytest.mark.asyncio
@@ -516,7 +516,7 @@ async def test_disable_params_not_propagated(self, decorator, decorator_call):
 
     @pytest.mark.asyncio
     async def test_set_in_cache(self, decorator, decorator_call):
-        await decorator.set_in_cache({"a": 1, "b": 2}, (), {})
+        await decorator.set_in_cache({"a": 1, "b": 2}, stub_dict, (), {})
 
         call_args = decorator.cache.multi_set.call_args[0][0]
         assert ("a", 1) in call_args
@@ -526,15 +526,15 @@ async def test_set_in_cache(self, decorator, decorator_call):
     @pytest.mark.asyncio
     async def test_set_in_cache_with_ttl(self, decorator, decorator_call):
         decorator.ttl = 10
-        await decorator.set_in_cache({"a": 1, "b": 2}, (), {})
+        await decorator.set_in_cache({"a": 1, "b": 2}, stub_dict, (), {})
 
         assert decorator.cache.multi_set.call_args[1]["ttl"] == decorator.ttl
 
     @pytest.mark.asyncio
     async def test_set_in_cache_exception(self, decorator, decorator_call):
         decorator.cache.multi_set = CoroutineMock(side_effect=Exception)
 
-        assert await decorator.set_in_cache({"a": 1, "b": 2}, (), {}) is None
+        assert await decorator.set_in_cache({"a": 1, "b": 2}, stub_dict, (), {}) is None
 
     @pytest.mark.asyncio
     async def test_decorate(self, mock_cache):