Search code examples
pythonpython-3.xdecoratorpython-decoratorspython-lru-cache

Python lru_cache nested decorator


I want to basically ignore a dict parameter in a method that want to decorate with functools.lru_cache.

This is the code I have written so far:

import functools


class BlackBox:
    """All BlackBoxes are the same."""
    def __init__(self, contents):
        self._contents = contents

    @property
    def contents(self):
        return self._contents

    def __eq__(self, other):
        return isinstance(other, type(self))

    def __hash__(self):
        return hash(type(self))

def lru_cache_ignore(func):
    lru_decorator = functools.lru_cache(maxsize=2048)

    @functools.wraps(func)
    def wrapper_decorator(self, arg1, *args, **kwargs):

        @lru_decorator
        def helper(arg1_w: BlackBox, *args, **kwargs):
            print(f"Cache miss {arg1_w} {args} {kwargs}")
            # unpack the first argument from Blackbox
            original_arg = arg1_w.contents
            # Invoke the user function seamlessly
            return func(self, original_arg, *args, **kwargs)

        boxed_arg = BlackBox(arg1)
        print(helper)
        print(helper.__wrapped__)
        return helper(boxed_arg, *args, **kwargs)

    return wrapper_decorator

The intention of the usage is clear in the unit tests as follow:

def test_actual_function_is_called():
    my_mock = MagicMock()

    class A:
        @lru_cache_ignore
        def my_fun(self, dict_arg, str_arg):
            my_mock.call_fun(dict_arg, str_arg)

    A().my_fun({}, 'test')
    my_mock.call_fun.assert_called_once_with({}, 'test')


def test_second_call_is_cached():
    my_mock = MagicMock()
    my_mock.call_fun.return_value = [1]

    class A:
        @lru_cache_ignore
        def my_fun(self, dict_arg, str_arg):
            return my_mock.call_fun(dict_arg, str_arg)
    a = A()
    first_result = a.my_fun({}, 'test')
    second_result = a.my_fun({}, 'test')

    my_mock.call_fun.assert_called_once_with({}, 'test')

    assert first_result == second_result

The issue is that the second test does not pass. The decorated function is called twice even though the input is the same.

This is the pytest output:

[CPython38-test] msg = "Expected 'call_fun' to be called once. Called 2 times.\nCalls: [call({}, 'test'), call({}, 'test')]."
[CPython38-test]
[CPython38-test]     def assert_called_once_with(_mock_self, *args, **kwargs):
[CPython38-test]         """assert that the mock was called exactly once and that that call was
[CPython38-test]         with the specified arguments."""
[CPython38-test]         self = _mock_self
[CPython38-test]         if not self.call_count == 1:
[CPython38-test]             msg = ("Expected '%s' to be called once. Called %s times.%s"
[CPython38-test]                    % (self._mock_name or 'mock',
[CPython38-test]                       self.call_count,
[CPython38-test]                       self._calls_repr()))
[CPython38-test] >           raise AssertionError(msg)
[CPython38-test] E           AssertionError: Expected 'call_fun' to be called once. Called 2 times.
[CPython38-test] E           Calls: [call({}, 'test'), call({}, 'test')].
[CPython38-test]
[CPython38-test] /.../mock.py:956: AssertionError
[CPython38-test] ----------------------------- Captured stdout call -----------------------------
[CPython38-test] <functools._lru_cache_wrapper object at 0x7f8eacf31430>
[CPython38-test] <function lru_cache_ignore.<locals>.wrapper_decorator.<locals>.helper at 0x7f8eacf27e50>
[CPython38-test] Cache miss <ads.adapter.rule_engine.lru_utils.BlackBox object at 0x7f8eacb7e8e0> ('test',) {}
[CPython38-test] <functools._lru_cache_wrapper object at 0x7f8eacf31430>
[CPython38-test] <function lru_cache_ignore.<locals>.wrapper_decorator.<locals>.helper at 0x7f8eacf27e50>
[CPython38-test] Cache miss <ads.adapter.rule_engine.lru_utils.BlackBox object at 0x7f8eacb7e8e0> ('test',) {}

I can see that the functions that are invoked in the decorator are the same. What am I missing here? I have adapted the answer from this other post

Thank you!


Solution

  • Issues

    1. The helper() function needs to be moved out of the wrapper_decorator. Otherwise, every call to the function creates a new helper function along with a brand new cache.

    2. I'm not sure of your intended semantics, but the self argument wasn't used consistently between the various levels of call forwarding.

    Fixed-up code

    
    def lru_cache_ignore(func):
    
        @functools.lru_cache(maxsize=2048)
        def helper(self, arg1_w: BlackBox, *args, **kwargs):
            print(f"Cache miss {arg1_w} {args} {kwargs}")
            # unpack the first argument from Blackbox
            original_arg = arg1_w.contents
            # Invoke the user function seamlessly
            return func(self, original_arg, *args, **kwargs)
    
        @functools.wraps(func)
        def wrapper_decorator(self, arg1, *args, **kwargs):
            boxed_arg = BlackBox(arg1)
            print(helper)
            print(helper.__wrapped__)
            return helper(self, boxed_arg, *args, **kwargs)
    
        return wrapper_decorator
    

    Comment

    The use of the BlackBox wrapper is ingenious and provides a clean way to pass in arguments that are unhashable or that should not be part of the cache key. I have not seen that before. Kudos.