Fix cache invalidation with custom key generator (#20)

src/cache_memoize/__init__.py

@@ -91,7 +91,7 @@ def noop(*args):
     cache = caches[cache_alias]
 
     def decorator(func):
-        def _make_cache_key(*args, **kwargs):
+        def _default_make_cache_key(*args, **kwargs):
             cache_key = ":".join(
                 [force_text(x) for x in args_rewrite(*args)]
                 + [force_text("{}={}".format(k, v)) for k, v in kwargs.items()]
@@ -100,16 +100,15 @@ def _make_cache_key(*args, **kwargs):
                 force_bytes("cache_memoize" + (prefix or func.__name__) + cache_key)
             ).hexdigest()
 
+        _make_cache_key = key_generator_callable or _default_make_cache_key
+
         @wraps(func)
         def inner(*args, **kwargs):
             # The cache key string should never be dependent on special keyword
             # arguments like _refresh. So extract it into a variable as soon as
             # possible.
             _refresh = bool(kwargs.pop("_refresh", False))
-            if key_generator_callable is None:
-                cache_key = _make_cache_key(*args, **kwargs)
-            else:
-                cache_key = key_generator_callable(*args, **kwargs)
+            cache_key = _make_cache_key(*args, **kwargs)
             if _refresh:
                 result = MARKER
             else:

tests/test_cache_memoize.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import random
 
 from django.core.cache import cache
 
@@ -170,8 +171,6 @@ def test_invalidate():
 
     calls_made = []
 
-    import random
-
     @cache_memoize(10)
     def function(argument):
         calls_made.append(argument)
@@ -192,8 +191,6 @@ def test_invalidate_with_refresh():
 
     calls_made = []
 
-    import random
-
     @cache_memoize(10)
     def function(argument):
         calls_made.append(argument)
@@ -255,6 +252,31 @@ def runmeonce(arg1, arg2):
     assert len(calls_made) == 2
 
 
+def test_invalidate_with_custom_key_generator():
+
+    calls_made = []
+
+    def key_generator(*args):
+        key = (":{}" * len(args)).format(*args)
+        return "custom_namespace:{}".format(key)
+
+    @cache_memoize(10, key_generator_callable=key_generator)
+    def runmeonce(arg1, arg2):
+        calls_made.append((arg1, arg2))
+        return arg1 + 1
+
+    runmeonce(1, 2)
+    runmeonce(1, 2)
+    assert len(calls_made) == 1
+    runmeonce.invalidate(999, 10)  # different args
+    assert runmeonce(1, 2)
+    assert len(calls_made) == 1
+
+    runmeonce.invalidate(1, 2)  # known args
+    assert runmeonce(1, 2)
+    assert len(calls_made) == 2
+
+
 def test_cache_memoize_none_value():
     calls_made = []