Python Decorator for Async and Sync Function without code duplication

Question

I have seen at least two examples of decorator that work with both normal def sync_func() and async def async_function().

But all these example basically duplicated the decorator code like this

import asyncio
import typing as t
import functools

R = t.TypeVar("R")
P = t.ParamSpec("P")


def decorator(func: t.Callable[P, R]) -> t.Callable[P, R]:
    if asyncio.iscoroutinefunction(func):
        @functools.wraps(func)
        async def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
            # … do something before
            result = await func(*args, **kwargs)
            # … do something after
            return result
    else:
        @functools.wraps(func)
        def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
            # … do something before [code duplication!]
            result = func(*args, **kwargs)
            # … do something after [code duplication!]
            return result
        
    return decorated

Is there a possibility to create a decorator without code duplication inside decorated ?

Answer 1

You can use the iterable property of Couroutine.send to run a async def function without using an event loop.

This inspired by a comment in cpython.

import asyncio
import typing as t
import functools
import time

R = t.TypeVar("R")
P = t.ParamSpec("P")


def is_coroutine(
    func: t.Callable[P, t.Any]
) -> t.TypeGuard[t.Callable[P, t.Coroutine[t.Any, t.Any, t.Any]]]:
    """Return true if the given func is a coroutine"""
    # partial call needs to be removed
    while isinstance(func, functools.partial):
        func = func.func
    if asyncio.iscoroutinefunction(func):
        return True
    if hasattr(func, "__call__"):
        # handle classes with `async def __call__(…)` as well
        func = func.__call__
        return asyncio.iscoroutinefunction(func)
    return False



def decorator(func: t.Callable[P, R]) -> t.Callable[P, R]:
    is_sync = not is_coroutine(func)

    @functools.wraps(func)
    async def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"{is_sync=}")
        print("Before")
        result = func(*args, **kwargs)
        if not is_sync:
            result = await t.cast(t.Awaitable[R], result)
        print(f"After with {result=}")
        return result

    if is_sync:
        @functools.wraps(func)
        def decorated_sync(*args: P.args, **kwargs: P.kwargs) -> R:
            """
            Execute coroutine by yield from all `await` calls without using event loop
            """
            try:
                decorated(*args, **kwargs).send(None)
            except StopIteration as result:
                return t.cast(R, result.value)
            raise RuntimeError("Never") # just to make mypy happy

        return decorated_sync

    return decorated  # type: ignore[return-value]

⚠️ Attention

Even so async def decorated(…) looks like a normal async function and you are able to await other async def functions, they can not access the event loop (anything from asyncio.* or functions that use it) if is_sync is True. You have been warned!

Best use only sync code within the decorator or use async code only after if not is_sync:.

Usage

It just works as expected

@decorator
def sync_call() -> str:
    time.sleep(0.1)
    return "sync"


@decorator
async def async_call() -> str:
    await asyncio.sleep(0.1)
    return "async"


async def run() -> None:
    res = sync_call()
    print(res)

    res = await async_call()
    print(res)


asyncio.run(run())

Leading to

is_sync=True
Before
After with result='sync'
sync
is_sync=False
Before
After with result='async'
async

Available on MyPy Play