Search code examples
pythoncachingdecoratorwrappercachetools

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.


Solution

  • 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