Context manager for caching a class method

lru_cache_class_method.py

class LruCacheClassMethod():
    def __init__(self, klass, method_name):
        self.klass = klass
        self.method_name = method_name
        self.original_method = getattr(self.klass, self.method_name)
    
    def __enter__(self):
        new_method = functools.lru_cache()(self.original_method)
        setattr(self.klass, self.method_name, new_method)
        return new_method

    def __exit__(self, _type, value, traceback):
        setattr(self.klass, self.method_name, self.original_method)

class MyClass:
    @classmethod
    def blop(truc: int):
        return truc + 1
        
with LruCacheClassMethod(MyClass, 'blop'):
    print(MyClass.blop(5))
    print(MyClass.blop(5))