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

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