In python,how to use decorator compatible with normal function and coroutine function simply?

Question

Here is my way,but I feel it is not very simple, any better way?

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

there is a lot of decorator, retry, add log etc,all will write this way,a bit ugly,right?

Answer 1

While there is no real problems with repeating the same code in specialized decorators. Here is how I'll approach a refactoring.

I will use a class decorator that keeps the accepts a pre-call function and a post-call function, both of which will be called with an instance of the decorator. The result of the pre-call function will be saved to an attribute of the decorator.

This is necessary for the special timing case where a delta needs to be computed.

I guess there may be other examples that may require the return value of a pre-call function execution.

I also save the result of the decorated function executed to the result attribute of the decorator instance. This allows the post call function to read this value for logging.

Here is an example implementation:

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

Timer Example

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

Running sync adder

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

Running async adder

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

Logging Example

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

Running sync subtractor:

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

Running async subtractor:

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