Modifying a cooldown decorator to work for methods instead of functions
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
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: 1foo: 3foo: 4bar: 1bar: 2bar: 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
This decorator works both with functions and methods, supports remaining
property and is implemented as a single class.
import timeclass cooldown: def __init__(self, timeout): self.timeout = timeout self.calltime = time.time() - timeout self.func = None self.obj = None def __call__(self, *args, **kwargs): if self.func is None: self.func = args[0] return self now = time.time() if now - self.calltime >= self.timeout: self.calltime = now if self.obj is None: return self.func.__call__(*args, **kwargs) else: return self.func.__get__(self.obj, self.objtype)(*args, **kwargs) def __get__(self, obj, objtype): self.obj = obj self.objtype = objtype return self @property def remaining(self): now = time.time() delta = now - self.calltime if delta >= self.timeout: return 0 return self.timeout - delta @remaining.setter def remaining(self, value): self.calltime = time.time() - self.timeout + value
# test with functions@cooldown(8)def test(*args): print('Function', *args)>>> test()Function>>> test()>>> test.remaining4.718205213546753>>> test.remaining = 0>>> test()Function
# test with methodsclass A: def __init__(self, value): self.value = value @cooldown(5) def a(self, *args): print('Method', self.value, *args)>>> a = A(7)>>> a.a()Method 7>>> a.a()>>> a.a.remaining3.589237892348223>>> a.a.remaining = 10>>> a.a(32)>>> a.a.remaining8.423482288923785>>> a.a.remaining = 0>>> a.a(32)Method 7 32