Python method decorator

memoized_method.py

# ref: https://stackoverflow.com/questions/33672412/python-functools-lru-cache-with-instance-methods-release-object

import functools
import weakref

def memoized_method(*lru_args, **lru_kwargs):
    def decorator(func):
        @functools.wraps(func)
        def wrapped_func(self, *args, **kwargs):
            # We're storing the wrapped method inside the instance. If we had
            # a strong reference to self the instance would never die.
            self_weak = weakref.ref(self)
            @functools.wraps(func)
            @functools.lru_cache(*lru_args, **lru_kwargs)
            def cached_method(*args, **kwargs):
                return func(self_weak(), *args, **kwargs)
            setattr(self, func.__name__, cached_method)
            return cached_method(*args, **kwargs)
        return wrapped_func
    return decorator

test_memoized_method.py

import unittest
import functools
import weakref

class TestMemoizedMethod(unittest.TestCase):
    def test_memoized_method(self):
        class MyClass:
            def __init__(self):
                self.counter = 0

            @memoized_method(maxsize=2)
            def my_method(self, arg):
                self.counter += 1
                return arg

        obj = MyClass()
        result1 = obj.my_method("hello")
        result2 = obj.my_method("hello")
        self.assertEqual(result1, "hello")
        self.assertEqual(result2, "hello")
        self.assertEqual(obj.counter, 1)
        
        result3 = obj.my_method("world")
        result4 = obj.my_method("world")
        self.assertEqual(result3, "world")
        self.assertEqual(result4, "world")
        self.assertEqual(obj.counter, 2)
        
        result5 = obj.my_method("hello")
        self.assertEqual(result5, "hello")
        self.assertEqual(obj.counter, 2)