修改冷却装饰器以用于方法而不是函数

Question

我正在尝试创建一个装饰器,它可以用于对它们应用"冷却"的方法,这意味着它们不能在一定的持续时间内被多次调用.我已经为函数创建了一个:

>>> @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

但是我需要这个来支持类的方法而不是普通的函数:

>>> 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

这是我创建的,使其适用于正常功能:

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

但这不适用于方法,因为它将_CooldownFunction对象传递为完全self忽略原始对象self.我如何使用方法,正确传递原始self而不是_CooldownFunction对象？

此外,用户需要能够动态更改剩余时间,这使得这更难(不能只是__get__用来返回functools.partial(self.__call__, obj)或其他东西):

>>> 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

编辑:它只需要为方法工作,而不是为方法和函数工作.

编辑2:这个设计有一个巨大的缺陷.虽然它适用于普通函数,我希望它分别装饰每个实例.目前,如果我有两个实例t1并且t2要打电话t1.f(),我无法再打电话,t2.f()因为冷却时间是f()方法而不是实例.我可能会为此使用某种字典,但在这种认识之后我更加迷失......

Answer 1

您可以重写类的__get__方法以使其成为描述符。当有人__get__从其包含对象中获取装饰方法并传递包含对象时，将调用该方法，然后您可以将其传递给原始方法。它返回一个实现您需要的功能的对象。

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

该Wrapper对象实现了__call__以及您想要的任何属性，因此将这些实现移至该对象中。它看起来像：

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