1 3
from functools import wraps
2

3 3
import hashlib
4 3
from urllib.parse import quote
5

6 3
from django.core.cache import caches, DEFAULT_CACHE_ALIAS
7

8 3
from django.utils.encoding import force_bytes
9

10 3
MARKER = object()
11

12

13 3
def cache_memoize(
14
    timeout,
15
    prefix=None,
16
    args_rewrite=None,
17
    hit_callable=None,
18
    miss_callable=None,
19
    key_generator_callable=None,
20
    store_result=True,
21
    cache_alias=DEFAULT_CACHE_ALIAS,
22
):
23
    """Decorator for memoizing function calls where we use the
24
    "local cache" to store the result.
25

26
    :arg int timeout: Number of seconds to store the result if not None
27
    :arg string prefix: If None becomes the function name.
28
    :arg function args_rewrite: Callable that rewrites the args first useful
29
    if your function needs nontrivial types but you know a simple way to
30
    re-represent them for the sake of the cache key.
31
    :arg function hit_callable: Gets executed if key was in cache.
32
    :arg function miss_callable: Gets executed if key was *not* in cache.
33
    :arg key_generator_callable: Custom cache key name generator.
34
    :arg bool store_result: If you know the result is not important, just
35
    that the cache blocked it from running repeatedly, set this to False.
36
    :arg string cache_alias: The cache alias to use; defaults to 'default'.
37

38
    Usage::
39

40
        @cache_memoize(
41
            300,  # 5 min
42
            args_rewrite=lambda user: user.email,
43
            hit_callable=lambda: print("Cache hit!"),
44
            miss_callable=lambda: print("Cache miss :("),
45
        )
46
        def hash_user_email(user):
47
            dk = hashlib.pbkdf2_hmac('sha256', user.email, b'salt', 100000)
48
            return binascii.hexlify(dk)
49

50
    Or, when you don't actually need the result, useful if you know it's not
51
    valuable to store the execution result::
52

53
        @cache_memoize(
54
            300,  # 5 min
55
            store_result=False,
56
        )
57
        def send_email(email):
58
            somelib.send(email, subject="You rock!", ...)
59

60
    Also, whatever you do where things get cached, you can undo that.
61
    For example::
62

63
        @cache_memoize(100)
64
        def callmeonce(arg1):
65
            print(arg1)
66

67
        callmeonce('peter')  # will print 'peter'
68
        callmeonce('peter')  # nothing printed
69
        callmeonce.invalidate('peter')
70
        callmeonce('peter')  # will print 'peter'
71

72
    Suppose you know for good reason you want to bypass the cache and
73
    really let the decorator let you through you can set one extra
74
    keyword argument called `_refresh`. For example::
75

76
        @cache_memoize(100)
77
        def callmeonce(arg1):
78
            print(arg1)
79

80
        callmeonce('peter')                 # will print 'peter'
81
        callmeonce('peter')                 # nothing printed
82
        callmeonce('peter', _refresh=True)  # will print 'peter'
83

84
    """
85

86 3
    if args_rewrite is None:
87

88 3
        def noop(*args):
89 3
            return args
90

91 3
        args_rewrite = noop
92

93 3
    def decorator(func):
94 3
        def _default_make_cache_key(*args, **kwargs):
95 3
            cache_key = ":".join(
96
                [quote(str(x)) for x in args_rewrite(*args)]
97
                + [quote("{}={}".format(k, v)) for k, v in kwargs.items()]
98
            )
99 3
            return hashlib.md5(
100
                force_bytes("cache_memoize" + (prefix or func.__qualname__) + cache_key)
101
            ).hexdigest()
102

103 3
        _make_cache_key = key_generator_callable or _default_make_cache_key
104

105 3
        @wraps(func)
106
        def inner(*args, **kwargs):
107
            # The cache backend is fetched here (not in the outer decorator scope)
108
            # to guarantee thread-safety at runtime.
109 3
            cache = caches[cache_alias]
110
            # The cache key string should never be dependent on special keyword
111
            # arguments like _refresh. So extract it into a variable as soon as
112
            # possible.
113 3
            _refresh = bool(kwargs.pop("_refresh", False))
114 3
            cache_key = _make_cache_key(*args, **kwargs)
115 3
            if _refresh:
116 3
                result = MARKER
117
            else:
118 3
                result = cache.get(cache_key, MARKER)
119 3
            if result is MARKER:
120 3
                result = func(*args, **kwargs)
121 3
                if not store_result:
122
                    # Then the result isn't valuable/important to store but
123
                    # we want to store something. Just to remember that
124
                    # it has be done.
125 3
                    cache.set(cache_key, True, timeout)
126
                else:
127 3
                    cache.set(cache_key, result, timeout)
128 3
                if miss_callable:
129 3
                    miss_callable(*args, **kwargs)
130 3
            elif hit_callable:
131 3
                hit_callable(*args, **kwargs)
132 3
            return result
133

134 3
        def invalidate(*args, **kwargs):
135
            # The cache backend is fetched here (not in the outer decorator scope)
136
            # to guarantee thread-safety at runtime.
137 3
            cache = caches[cache_alias]
138 3
            kwargs.pop("_refresh", None)
139 3
            cache_key = _make_cache_key(*args, **kwargs)
140 3
            cache.delete(cache_key)
141

142 3
        inner.invalidate = invalidate
143 3
        return inner
144

145 3
    return decorator

Read our documentation on viewing source code .

Loading