Search code examples
pythonpython-3.xdecoratorpython-decoratorsfunctools

How to use these two class decorators together with functools.update_wrapper?


I was reading about decorators and tried to mix these two examples and make them class decorators instead of regular functions. The first one only lets you run a function once per argument and the second one counts how many times you've run that function. They both work fine separated but when I try to decorate a simple function with both at the same time it fails... Or doesn't really fail but prints an unexpected wrong result. I did some reading and found that the functools module can help but I'm not sure how.

from functools import update_wrapper

class Memoize:
    def __init__(self, func):
        self.func = func
        self.memo = dict()
        update_wrapper(self, func)

    def __call__(self, *args):
        if args not in self.memo:
            self.memo[args] = self.func(args)
        else:
            print("cls decorator. You have printed this before")
        return self.memo[args]


class CallCounter:
    def __init__(self, func):
        self.func = func
        self.calls = 0
        self.__name__ = func.__name__
        update_wrapper(self, func)

    def __call__(self, *args, **kwargs):
        self.calls += 1
        return self.func(*args, **kwargs)


@Memoize
@CallCounter
def doubleprint(x):
    for elem in x:
        print(elem + " " + elem)


doubleprint('Hello')
doubleprint('Hello')
doubleprint('Hello')
doubleprint('Hello')
doubleprint('Bye')

print(doubleprint.calls)

doubleprint('Hello')

doubleprint('Hello')
doubleprint('Hello')
doubleprint('Hello')
doubleprint('Bye')

print(doubleprint.calls)

Solution

  • By default update_wrapper updates the __dict__ from a wrapped class. So your func in Memoize is being overwritten with the func in CallCounter which means Memoize is directly calling your doubleprint() function and never calling CallCounter.

    class Memoize:
        def __init__(self, func):
            self.func = func
            print(type(self.func))  # <class '__main__.CallCounter'>
            self.memo = dict()
            update_wrapper(self, func)
            print(type(self.func))  # <class 'function'>
    

    You can fix this by doing:

            update_wrapper(self, func, updated=[])
    

    Which will not copy the __dict__ from the CallCounter into the Memoize instance but will still copy __name__, __doc__, etc.

    To access CallCounter class you would:

    print(doubleprint.__wrapped__.calls)
    

    But you need the fix above or this will always print 0 (because it is never called).