Modifying a cooldown decorator to work for methods instead of functions

Question

I'm trying to create a decorator which would work for methods to apply a "cooldown" on them, meaning they can't be called multiple times within a certain duration. I already created one for functions:

>>> @cooldown(5)
... def f():
...     print('f() was called')
...
>>> f()
f() was called
>>> f()  # Nothing happens when called immediately
>>> f()  # This is 5 seconds after first call
f() was called

but I need this to support methods of classes instead of normal functions:

>>> class Test:
...    @cooldown(6)
...    def f(self, arg):
...        print(self, arg)
...
>>> t = Test()
>>> t.f(1)
<Test object at ...> 1
>>> t.f(2)
>>> t.f(5)  # Later
<Test object at ...> 5

Here's what I created to make it work for normal functions:

import time

class _CooldownFunc:
    def __init__(self, func, duration):
        self._func = func
        self.duration = duration
        self._start_time = 0

    @property
    def remaining(self):
        return self.duration - (time.time() - self._start_time)

    @remaining.setter
    def remaining(self, value):
        self._start_time = time.time() - (self.duration - value)

    def __call__(self, *args, **kwargs):
        if self.remaining <= 0:
            self.remaining = self.duration
            return self._func(*args, **kwargs)

    def __getattr__(self, attr):
        return self._func.__getattribute__(attr)


def cooldown(duration):
    def decorator(func):
        return _CooldownFunc(func, duration)
    return decorator

But this doesn't work with methods, since it passes the _CooldownFunction object as self and ignores the original self completely. How would I get it to work with methods, properly passing the original self instead of the _CooldownFunction object?

Also, it is required for users to be able to change the remaining time on the fly, which makes this even harder (can't just use __get__ to return functools.partial(self.__call__, obj) or something):

>>> class Test:
...     @cooldown(10)
...     def f(self, arg): 
...         print(self, arg)
...
>>> t = Test()
>>> t.f(5)
<Test object at ...> 5
>>> t.f.remaining = 0
>>> t.f(3)  # Almost immediately after previous call
<Test object at ...> 3

Edit: It only needs to work for methods, not for both methods and functions.

Edit 2: There's a huge flaw in this design to begin with. While it works just fine for normal functions, I want it to decorate each instance separately. Currently if I were to have two instances t1 and t2 and were to call t1.f(), I could not longer call t2.f() because the cooldown is teid to the f() method instead of the instances. I probably could use some kind of dictionary for this, but after this realization I'm even more lost...

Answer 1

You can override your class's __get__ method to make it a descriptor. The __get__ method will be called when someone gets the decorated method from within its containing object, and is passed the containing object, which you will then be able to pass to the original method. It returns an object which implements the functionality you need.

def __get__(self, obj, objtype):
    return Wrapper(self, obj)

The Wrapper object implements __call__, and any properties you want, so move those implementations into that object. It would look like:

class Wrapper:
    def __init__(self, cdfunc, obj):
        self.cdfunc = cdfunc
        self.obj = obj
    def __call__(self, *args, **kwargs):
        #do stuff...
        self.cdfunc._func(self.obj, *args, **kwargs)
    @property
    def remaining(self):
        #...get needed things from self.cdfunc

Answer 2

Fixing the issue interjay addressed, I did a quick re-write of your cooldown decorator, which now works for all kinds of functions/methods:

class cooldown(object):
    def __init__(self, duration):
        self._duration = duration
        self._storage = self
        self._start_time = 0

    def __getRemaining(self):
        if not hasattr(self._storage, "_start_time"):
            self._storage._start_time = 0
        return self._duration - (time.time() -
                                 self._storage._start_time)

    def __setRemaining(self, value):
        self._storage._start_time = time.time() - (self._duration -
                                                   value)

    remaining = property(__getRemaining, __setRemaining)

    def __call__(self, func):
        is_method = inspect.getargspec(func).args[0] == 'self'
        def call_if(*args, **kwargs):
            if is_method :
                self._storage = args[0]
            else:
                self._storage = self
            if self.remaining <= 0:
                self.remaining = self._duration
                return func(*args, **kwargs)

        call_if.setRemaining = self.__setRemaining
        call_if.getRemaining = self.__getRemaining
        return call_if

Tests:

@cooldown(2)
def foo(stuff):
    print("foo: %s" % stuff)

foo(1)
foo(2)
time.sleep(3)
foo(3)
foo.setRemaining(0)
foo(4)

class Bla(object):
    @cooldown(2)
    def bar(self, stuff):
        print("bar: %s" % stuff)

bla = Bla()
bla.bar(1)
bla.bar.setRemaining(0)
bla.bar(2)
time.sleep(3)
bla.bar(3)
bla.bar(4)

outputs:

foo: 1
foo: 3
foo: 4
bar: 1
bar: 2
bar: 3

EDIT: I altered the code so it works independently for multiple instances by putting it's storage into the invoked function's self argument. Please note that this purely relies on the first argument being named "self", but you can search for a more robust way of detecting if a decorated callable is a method or a function if you need more security here.

EDIT2: This might has a bug if you do instance1.foo() and then try to do instance2.foo.setRemaining(0). Since the context didn't get switched, this would set the remaining value for instance1. Could be fixed by making the setters and getters bound methods to the context, but this is getting messy. I will stop here for now