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

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

The other try was through @typing.overload 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.