在python中，如何简单地使用兼容普通函数和协程函数的装饰器？

Question

这是我的方法，但我觉得不是很简单，有什么更好的方法吗？

import asyncio
import time


def timer_all(f):
    if asyncio.iscoroutinefunction(f):
        async def wrapper(*args, **kwargs):
            now = time.time()
            result = await f(*args, **kwargs)
            print('used {}'.format(time.time() - now))
            return result
    else:
        def wrapper(*args, **kwargs):
            now = time.time()
            result = f(*args, **kwargs)
            print('used {}'.format(time.time() - now))
            return result

    return wrapper

有很多装饰器，重试，添加日志等，都会这样写，有点丑，对吧？

Answer 1

虽然在专门的装饰器中重复相同的代码并没有真正的问题。以下是我将如何进行重构。

我将使用一个类装饰器来保留接受预调用函数和后调用函数，这两个函数都将使用装饰器的实例进行调用。预调用函数的结果将保存到装饰器的属性中。

这对于需要计算增量的特殊时序情况是必要的。

我想可能还有其他示例可能需要预调用函数执行的返回值。

我还将执行的装饰函数的结果保存到装饰器实例的 result 属性中。这允许后调用函数读取该值以进行记录。

这是一个示例实现：

import asyncio


class WrapAll(object):
    def __init__(self, pre=lambda _: None, post=lambda _: None):
        self.pre = lambda : pre(self)
        self.pre_val = None
        self.result = None
        self.post = lambda : post(self)

    def __call__(self, fn):
        if asyncio.iscoroutinefunction(fn):
            async def wrap(*args, **kwargs):
                self.pre_val = self.pre()
                self.result = await fn(*args, *kwargs)
                self.post()
                return self.result
        else:
            def wrap(*args, **kwargs):
                self.pre_val = self.pre()
                self.result = fn(*args, *kwargs)
                self.post()
                return self.result
        return wrap

定时器示例

import asyncio
import time

timer = dict(
    pre=lambda self: time.time(),
    post=lambda self: print('used {}'.format(time.time()-self.pre_val))
)

@WrapAll(**timer)
def add(x, y):
    return x + y

@WrapAll(**timer)
async def async_add(x, y):
    future = asyncio.Future()
    future.set_result(x+y)
    await future
    return future.result()

运行同步加法器

>>> add(3, 4)
used 4.76837158203125e-06
7

运行异步加法器

>>> loop = asyncio.get_event_loop()
>>> task = asyncio.ensure_future(async_add(3, 4))
>>> try:
...    loop.run_until_complete(task)
... except RuntimeError:
...     pass
used 2.193450927734375e-05

记录示例

import asyncio
import logging

FORMAT = '%(message)s'
logging.basicConfig(format=FORMAT)

logger = dict(
    post=lambda self: logging.warning('subtracting {}'.format(self.result))
)

@WrapAll(**logger)
def sub(x, y):
    return x - y

@WrapAll(**logger)
async def async_sub(x, y):
    future = asyncio.Future()
    future.set_result(x-y)
    await future
    return future.result()

运行同步减法器：

>>> sub(5, 6)
subtracting -1

运行异步减法器：

>>> loop = asyncio.get_event_loop()
>>> task = asyncio.ensure_future(async_sub(5, 6))
>>> try:
...     loop.run_until_complete(task)
... except RuntimeError:
...     pass
subtracting -1