Wrapping a decorator from an external library with arguments

I'm using the cachetools library and I would like to wrap the decorator form this library and add a class self argument to enable/disable the caching at the class level e.e. MyClass(enable_cache=True)

An example usage would be something like:

class MyClass(object):
    def __init__(self, enable_cache=True):
        self.enable_cache = enable_cache
        self.cache = cachetools.LRUCache(maxsize=10)
    
    @cachetools.cachedmethod(operator.attrgetter('cache'))
    def calc(self, n):
        return 1*n

I'm not sure how to keep the cache as a shared self class object and allow for the enable_cache flag within my own wrapper decorator using this library.

If I understand you correctly the answer is actually quite simple - you set the cache to None.

import cachetools
import operator

class MyClass(object):
    def __init__(self, enable_cache=True):
        self.cache = cachetools.LRUCache(maxsize=10) if enable_cache else None
    
    @cachetools.cachedmethod(operator.attrgetter('cache'))
    def calc(self, n):
        print("Calculating", n)
        return 1*n
    
    
m1 = MyClass(True)
m1.calc(2)
m1.calc(2)
m1.calc(3)
m1.calc(3)
print("now without")
m2 = MyClass(False)
m2.calc(2)
m2.calc(2)
m2.calc(3)
m2.calc(3)

Output:

Calculating 2
Calculating 3
now without
Calculating 2
Calculating 2
Calculating 3
Calculating 3

More flexible you can do it do it by wrapping the cache or by making a whole new decorator:

import cachetools
import operator

def flexible_cache(cache):
    def cache_wrapper(self):
        if self.enable_cache:
            return cache(self)
        return None
    return cache_wrapper

def optional_cache(cache, *args, **kwargs):
    return cachetools.cachedmethod(flexible_cache(cache), *args, **kwargs)
    

class MyClass(object):
    def __init__(self, enable_cache=True):
        self.enable_cache = enable_cache
        self.cache = cachetools.LRUCache(maxsize=10) # Now the None part is handled by the decorators
    
    @cachetools.cachedmethod(flexible_cache(operator.attrgetter('cache')))
    def calc2(self, n):
        print("Calculating2", 2*n)
        return 2*n
    
    @optional_cache(operator.attrgetter('cache'))
    def calc3(self, n):
       print("Calculating3", 2*n)
       return 2*n 

I understand that you want to defer the choice to decorate to the user who will instantiate your class. Then setattr is your friend. I demonstrate its use with a custom decorator, but the principle is the same.

def decorator(func):
    def inner(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f"{result = }")
        return result
    return inner

class Asd:
    def __init__(self, decorate=False):
        if decorate:
            setattr(self, "method", decorator(getattr(self, "method")))
    
    def method(self):
        return 42

>>> asd = Asd(decorate=False)
>>> asd.method()
42
>>> asd = Asd(decorate=True)
>>> asd.method()
result = 42
42

To achieve the desired functionality of enabling or disabling caching at the class level using the cachetools library, you can create a custom decorator that wraps the cachedmethod decorator. Here's an example implementation:

import cachetools
import operator

def class_cachedmethod(cache_key, maxsize=128):
    def decorator(cls):
        cls.cache = cachetools.LRUCache(maxsize=maxsize)
        
        def wrapper(method):
            if not getattr(cls, 'enable_cache', True):
                return method
            return cachetools.cachedmethod(operator.attrgetter(cache_key))(method)
        
        setattr(cls, cache_key, wrapper)
        return cls
    return decorator

In the code above, we define a class_cachedmethod decorator that takes a cache_key argument, which represents the cache attribute name within the class. The decorator returns another decorator that wraps the class and its methods.

Here's how you can use it with your example:

@class_cachedmethod('cache')
class MyClass(object):
    def __init__(self, enable_cache=True):
        self.enable_cache = enable_cache

    def cache(self, method):
        return method

    @cache
    def calc(self, n):
        return 1 * n

In this example, we apply the class_cachedmethod decorator to the MyClass class, specifying 'cache' as the cache attribute name. The calc method is decorated with the @cache decorator, which internally checks the enable_cache flag and decides whether to apply caching or not.

If enable_cache is True, the calc method will be cached using the cachetools.cachedmethod decorator with the cache attribute operator.attrgetter('cache'). If enable_cache is False, the calc method will be returned without caching.

By default, the cache size is set to 128, but you can adjust it by modifying the maxsize parameter in the class_cachedmethod decorator or in the LRUCache instantiation inside the decorator.