Type inference fails on decorator designed to work on both async and regular defs

[adampauls]

Describe the bug
I'm trying to write a decorator that takes an argument and decorates either an async or regular def in the case where there is a shared type parameter in the input and output of the def. Unlike the case discussed in #2142, the two overloads differ only in their outputs. I am aware that in the case the problem can be fixed by using the decorator as @foo instead of @foo(), but in the case we are worried about, @foo() takes a non-empty list of parameters.

To Reproduce

Output = TypeVar("Output")

@overload
def foo() -> (
    Callable[[Callable[[Output], Coroutine[Any, Any, Output]]], Callable[[Output], Coroutine[Any, Any, Output]]]
):
    ...


@overload
def foo() -> Callable[[Callable[[Output], Output]], Callable[[Output], Output]]:
    ...


def foo() -> (
    Callable[[Callable[[Output], Coroutine[Any, Any, Output]]], Callable[[Output], Coroutine[Any, Any, Output]]]
    | Callable[[Callable[[Output], Output]], Callable[[Output], Output]]
):
    return lambda x: x  # do something real here depending on value of `iscoroutinefunction`


def test_foo():
    @foo()
    def bar(x: int) -> int:
        return x + 5

    assert bar(1) == 6

    @foo()
    async def abar(x: int) -> int:
        return x + 5

    assert asyncio.run(abar(1)) == 6

Expected behavior
I think Pyright should successfully typecheck here.

The current behavior depends on the order in which the two @overloads are declared, which indicates that the output constraint on the decorator invocation from def's signature is not active when running type inference for the decorator.

VS Code extension or command-line
Command-line with 1.1.304.

[erictraut]

Pyright's behavior is correct here.

You have a bug in your overload definition.

When you invoke foo(), you're invoking an overloaded function. Overloads are selected based on the input arguments, but in this case there is no input argument because both overloads have no parameters in their signature. That means the first of the two overloads If you enable the reportOverlappingOverload, pyright will notify you over bugs like this.

[adampauls]

Just to double-check: there is no workaround then? There is no way to define a typesafe decorator that takes arguments and can work on both an async and non-async def?

[erictraut]

It's possible, but you need to move the @overload into a callback protocol class. This is necessary because Callable is very limited and can't be used to specify an overload.

class Decorator(Protocol):
    @overload
    def __call__(
        self, func: Callable[[Output], Coroutine[Any, Any, Output]]
    ) -> Callable[[Output], Coroutine[Any, Any, Output]]:
        ...

    @overload
    def __call__(self, func: Callable[[Output], Output]) -> Callable[[Output], Output]:
        ...


def foo() -> Decorator:
    ...

[adampauls]

Thanks, we tried something along those lines but couldn't get it to work. Will try again!