Annotating the decorator with arguments that works for both sync/async functions

[Galtozzy]

Hey there!

For context - I've tried asking the same question on typing Gitter and had no luck there since I believe the visibility is quite low and I got no comments to the message up to now.
I've checked the issues 1, 2

I am taking my luck here just to understand if what I'm trying to achieve is possible at all or if I shouldn't bother trying.

The idea seems simple - we need to preserve all types of a function the decorator with arguments is being applied to. However, the issue comes when the decorator is supposed to work with both sync and async functions.

The solution I've come close with is the following:

R = TypeVar('R')
P = ParamSpec('P')


def is_coroutine(func: Union[Callable[P, R], Callable[P, Awaitable[R]]]) -> TypeGuard[Callable[P, Awaitable[R]]]:
    return inspect.iscoroutinefunction(func)


def my_dec(
    param1: str, param2: Union[int, None] = None
) -> Callable[[Union[Callable[P, Awaitable[R]], Callable[P, R]]], Union[Callable[P, Awaitable[R]], Callable[P, R]]]:

    def decorator(_func: Union[Callable[P, R], Callable[P, Awaitable[R]]]) -> Union[Callable[P, R], Callable[P, Awaitable[R]]]:
        if is_coroutine(_func):
            _awaitable_func = _func

            @wraps(_awaitable_func)
            async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                return await _awaitable_func(*args, **kwargs)
            return _async_wrapper

        else:

            @wraps(_func)
            def _sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                return cast(R, _func(*args, **kwargs))

            return cast(Callable[P, R], _sync_wrapper)

    return decorator
    
    
@my_dec(param1='test')
def test() -> str:
    return 'test return'
    
    
return_value: str = test()

However, TypeGuard doesn't narrow down the type thus mypy throws an error on return_value: str = test():
Incompatible types in assignment (expression has type "Awaitable[str] | str", variable has type "str") [assignment]

I have also tried @typing.overload uproach but it seems it doesn't work in this case since overloading happens inside the my_dec.
I would appreciate any help, directions and resources I could look into to solve this.

[erictraut]

Here's a potential solution. It type checks without errors in pyright. Mypy produces a few overload-related errors that appear to be false positives, so you could add # type: ignore comments to suppress them.

Code sample in pyright playground

from typing import Awaitable, Callable, Protocol, overload
from typing_extensions import TypeIs
from inspect import iscoroutinefunction
from functools import wraps

def is_coroutine[**P, R](
    func: Callable[P, R | Awaitable[R]],
) -> TypeIs[Callable[P, Awaitable[R]]]:
    return iscoroutinefunction(func)

class SyncOrAsync(Protocol):
    @overload
    def __call__[**P, R](
        self, _func: Callable[P, Awaitable[R]]
    ) -> Callable[P, Awaitable[R]]:
        ...

    @overload
    def __call__[**P, R](self, _func: Callable[P, R]) -> Callable[P, R]:
        ...

    def __call__[**P, R](
        self, _func: Callable[P, Awaitable[R]] | Callable[P, R]
    ) -> Callable[P, Awaitable[R]] | Callable[P, R]:
        ...

def my_dec(param1: str, param2: int | None = None) -> SyncOrAsync:
    @overload
    def decorator[**P, R](
        _func: Callable[P, Awaitable[R]],
    ) -> Callable[P, Awaitable[R]]:
        ...

    @overload
    def decorator[**P, R](
        _func: Callable[P, R],
    ) -> Callable[P, R]:
        ...

    def decorator[**P, R](
        _func: Callable[P, Awaitable[R]] | Callable[P, R],
    ) -> Callable[P, Awaitable[R]] | Callable[P, R]:
        if is_coroutine(_func):
            _awaitable_func = _func

            @wraps(_awaitable_func)
            async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                return await _awaitable_func(*args, **kwargs)

            return _async_wrapper

        else:

            @wraps(_func)
            def _sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                return _func(*args, **kwargs)

            return _sync_wrapper

    return decorator

@my_dec(param1="test")
def test_sync() -> str:
    return "test return"

v_sync = test_sync()

@my_dec(param1="test")
async def test_async() -> str:
    return "test return"

v_async = test_async()

[Galtozzy]

Worked like a charm, thank you very much!