Optional caching in python - Wrapping a cachetools decorator 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.

When you use cachetools 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) # print
m1.calc(2) # should not print
m1.calc(3) # print
m1.calc(3) # should not print
print("now without")
m2 = MyClass(False)
m2.calc(2) # print
m2.calc(2) # print
m2.calc(3) # print
m2.calc(3) # print

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