类作为 Python 中常规函数和协程的装饰器

Question

我正在尝试创建一个类作为装饰器，它将一个 try-except 块应用于装饰函数并保留一些异常日志。我想将装饰器应用于常规函数和协程。

我已经完成了 class-as-decorator，它的工作原理与为常规函数设计的一样，但是协程出了点问题。以下是类作为装饰器的简化版本和几个用例的一些简化代码：

import traceback
import asyncio
import functools

class Try:

    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f"applying __call__ to {self.func.__name__}")
        try:
            return self.func(*args, **kwargs)
        except:
            print(f"{self.func.__name__} failed")
            print(traceback.format_exc())

    def __await__(self, *args, **kwargs):
        print(f"applying __await__ to {self.func.__name__}")
        try:
            yield self.func(*args, **kwargs)
        except:
            print(f"{self.func.__name__} failed")
            print(traceback.format_exc())

# Case 1
@Try
def times2(x):
    return x*2/0

# Case 2
@Try
async def times3(x):
    await asyncio.sleep(0.0001)
    return x*3/0

async def test_try():
    return await times3(10)

def main():
    times2(10)
    asyncio.run(test_try())
    print("All done")

if __name__ == "__main__":
    main()

这是上述代码的输出（稍作修改）：

applying __call__ to times2
times2 failed
Traceback (most recent call last):
  File "<ipython-input-3-37071526b2e6>", line 14, in __call__
    return self.func(*args, **kwargs)
  File "<ipython-input-3-37071526b2e6>", line 30, in times2
    return x*2/0
ZeroDivisionError: division by zero

applying __call__ to times3
Traceback (most recent call last):
  File "[...]/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3296, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-3-37071526b2e6>", line 46, in <module>
    main()
  File "<ipython-input-3-37071526b2e6>", line 43, in main
    asyncio.run(test_try())
  File "[...]/lib/python3.7/asyncio/runners.py", line 43, in run
    return loop.run_until_complete(main)
  File "[...]/lib/python3.7/asyncio/base_events.py", line 579, in run_until_complete
    return future.result()
  File "<ipython-input-3-37071526b2e6>", line 39, in test_try
    return await times3(10)
  File "<ipython-input-3-37071526b2e6>", line 36, in times3
    return x*3/0
ZeroDivisionError: division by zero

情况 1 表现正常：正如预期的那样，__call__ 被调用，然后装饰函数失败并捕获异常。但我无法解释案例 2 的行为。请注意最后缺少的“times3 failed”和“All done”打印。我无法在这里重现颜色编码的输出，但案例 1 的回溯是常规打印，而案例 2 的回溯是异常红色（在 PyCharm 上）。令人惊讶的是，调用了 __call__ 方法而不是 __await__。

我尝试了另一个类作为装饰器，它记录了函数被调用的次数。这与 __call__ 与常规函数或协程一起工作得很好。

那么实际上发生了什么？我是否需要以某种方式强制该函数使用__await__？怎么样？

我尝试了以下方法：

async def test_try2():
    func = await times3

有输出

applying __await__ to times3
times3 failed
Traceback (most recent call last):
  File "<ipython-input-5-5a85f988097e>", line 22, in __await__
    yield self.func(*args, **kwargs)
TypeError: times3() missing 1 required positional argument: 'x'

这会强制使用__await__，但然后呢？

Answer 1

您的代码的问题在于它将__await__ 放置在错误的对象上。通常await f(x) 会扩展为：

_awaitable = f(x)
_iter = _awaitable.__await__()
yield from _iter  # not literally[1]

注意__await__() 是如何在函数的result 上调用的，而不是在函数对象本身上。在您的 times3 示例中发生的情况如下：

• __call__ 调用self.func 中的原始times3 协程函数，它简单地构造了一个协程对象。此时也不例外，因为对象还没有开始执行，所以返回了一个coroutine object（调用async def协程函数得到的）。

• __await__ 是在通过运行self.func 获得的协程对象上调用的，它是原始的times3 async def，而不是在您的函数包装器上。这是因为，就上面的伪代码而言，您的包装器对应于f，而__await__() 是在_awaitable 上调用的，在您的情况下，这是调用f 的结果。

  李>

一般来说，您无法知道函数调用的结果是否会被等待。但是由于协程对象除了等待它们之外没有任何用处（它们甚至在没有等待的情况下被销毁时会打印警告），因此您可以放心地假设。这个假设允许您的 __call__ 检查函数调用的结果是否可等待，如果是，则将其包装在一个对象中，该对象将在 __await__ 级别实现包装逻辑：

...
import collections.abc

class Try:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f"applying __call__ to {self.func.__name__}")
        try:
            result = self.func(*args, **kwargs)
        except:
            print(f"{self.func.__name__} failed")
            print(traceback.format_exc())
            return
        if isinstance(result, collections.abc.Awaitable):
            # The result is awaitable, wrap it in an object
            # whose __await__ will call result.__await__()
            # and catch the exceptions.
            return TryAwaitable(result)
        return result

class TryAwaitable:
    def __init__(self, awaitable):
        self.awaitable = awaitable

    def __await__(self, *args, **kwargs):
        print(f"applying __await__ to {self.awaitable.__name__}")
        try:
            return yield from self.awaitable.__await__()
        except:
            print(f"{self.awaitable.__name__} failed")
            print(traceback.format_exc())

这会产生预期的输出：

applying __call__ to times3
applying __await__ to times3
times3 failed
Traceback (most recent call last):
  File "wrap3.py", line 30, in __await__
    yield from self.awaitable.__await__()
  File "wrap3.py", line 44, in times3
    return x*3/0
ZeroDivisionError: division by zero

请注意，__await__ 的实现有一个不相关的问题，它使用yield 委托给函数。必须改用yield from，因为这允许底层迭代器选择何时暂停，并在停止暂停时提供一个值。裸露的yield 无条件挂起（且仅挂起一次），这与await 的语义不兼容。

¹不是字面意思，因为async def 中不允许使用yield from。但是async def 的行为就像这样一个生成器是由它返回的对象的__await__ 方法返回的。