September 1, 2017

Very useful redis cache decorator in Python

Lets start from the end.

@cache('1h')
def func1(user_id):
    return heavy_compute(user_id)

result = func1(123)

Now our decorated func1 has some very useful methods!

To get value strictly from redis (do not compute, even if it absent in redis) we can use .get() method:

result_from_redis = func1.get(123)

To set (update if already exists) value to redis we can use .update() method. It will recompute value and save it in redis:

newly_computed_result = func1.update(123)

To save in redis value we already have (not to compute it) we can use .update_manually() method. It accepts func1 args + value to save (last arg):

func1.update_manually(123, 'winner')  
# 123: user_id
# 'winner': our result to save in redis

To delete result from redis we can use .delete() method:

deleted_count = func1.delete(123)

To get redis key we can use .get_key() method:

key_in_redis = func1.get_key(123)

Here full code for this cache decorator:

import json
from functools import partial, wraps


redis = Redis()  # for example Flask-Redis class
null = object()


def _get_cache_key(func, *args):
    result = '%s.%s:' % (func.__module__, func.__name__)
    if args:
        result = '%s%s' % (result, str(args))

    return result


def _get_cache(func, *args):
    value = redis.client.get(_get_cache_key(func, *args))
    if value is not None:
        return json.loads(value)

    return null


def _update_cache(func, *args):
    value = func(*args)
    redis.client.set(
        name=_get_cache_key(func, *args), 
        value=json.dumps(value), 
        ex=func.timeout,
    )

    return value


def _update_cache_manually(func, *args_with_value_as_last):
    redis.client.set(
        name=_get_cache_key(func, *args_with_value_as_last[:-1]),
        value=json.dumps(args_with_value_as_last[-1]),
        ex=func.timeout,
    )

    return args_with_value_as_last[-1]


def _delete_cache(func, *args):
    return redis.client.delete(_get_cache_key(func, *args))


_MAP = dict(
    s=1,  # seconds
    m=60,  # minutes
    h=60 * 60,  # hours
    d=60 * 60 * 24,  # days
    w=60 * 60 * 24 * 7,  # weeks
)


def cache(timeout):
    """@timeout: 60, '40s', '5m', '1h', '2d', '3w'"""
    def wrapper(func):
        if isinstance(timeout, int):
            seconds = timeout
        elif timeout[-1] in _MAP:  # last letter [s, m, h, d, w]
            seconds = int(timeout.rstrip(timeout[-1])) * _MAP[timeout[-1]]
        else:
            raise ValueError('Unknown format for timeout `%s`' % timeout)

        func.timeout = seconds

        @wraps(func)
        def func_wrapped(*args):
            value = _get_cache(func, *args)
            if value is not null:
                return value

            return _update_cache(func, *args)

        func_wrapped.timeout = func.timeout
        func_wrapped.get_key = partial(_get_cache_key, func)
        func_wrapped.get = partial(_get_cache, func)
        func_wrapped.update = partial(_update_cache, func)
        func_wrapped.update_manually = partial(_update_cache_manually, func)
        func_wrapped.delete = partial(_delete_cache, func)

        return func_wrapped

    return wrapper

No comments:

Post a Comment