How to create a Python decorator that can wrap either coroutine or function?

Question

I am trying to make a decorator to wrap either coroutines or functions.

The first thing I tried was a simple duplicate code in wrappers:

def duration(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_ts = time.time()
        result = func(*args, **kwargs)
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))
        return result

    @functools.wraps(func)
    async def async_wrapper(*args, **kwargs):
        start_ts = time.time()
        result = await func(*args, **kwargs)
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))
        return result

    if asyncio.iscoroutinefunction(func):
        return async_wrapper
    else:
        return wrapper

This works, but i want to avoid duplication of code, as this is not much better than writing two separate decorators.

Then i tried to make a decorator using class:

class SyncAsyncDuration:
    def __init__(self):
        self.start_ts = None

    def __call__(self, func):
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            self.setup(func, args, kwargs)
            result = func(*args, **kwargs)
            self.teardown(func, args, kwargs)
            return result

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            self.setup(func, args, kwargs)
            result = await func(*args, **kwargs)
            self.teardown(func, args, kwargs)
            return result

        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper

    def setup(self, func, args, kwargs):
        self.start_ts = time.time()

    def teardown(self, func, args, kwargs):
        dur = time.time() - self.start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

That works in some cases very well for me, but in this solution i can't put a function in with or try statements. Is there any way i can create a decorator without duplicating code?

Answer 1

May be you can find better way to do it, but, for example, you can just move your wrapping logic to some context manager to prevent code duplication:

import asyncio
import functools
import time
from contextlib import contextmanager


def duration(func):
    @contextmanager
    def wrapping_logic():
        start_ts = time.time()
        yield
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if not asyncio.iscoroutinefunction(func):
            with wrapping_logic():
                return func(*args, **kwargs)
        else:
            async def tmp():
                with wrapping_logic():
                    return (await func(*args, **kwargs))
            return tmp()
    return wrapper